diff --git a/needle/v1/__init__.py b/needle/v1/__init__.py index 79ac192..96848aa 100644 --- a/needle/v1/__init__.py +++ b/needle/v1/__init__.py @@ -11,6 +11,7 @@ NeedleBaseClient, ) from needle.v1.collections import NeedleCollections +from needle.v1.files import NeedleFiles NEEDLE_DEFAULT_URL = "https://needle-ai.com" @@ -46,3 +47,4 @@ def __init__( # sub clients self.collections = NeedleCollections(config, headers) + self.files = NeedleFiles(config, headers) diff --git a/needle/v1/files/__init__.py b/needle/v1/files/__init__.py new file mode 100644 index 0000000..b9882d8 --- /dev/null +++ b/needle/v1/files/__init__.py @@ -0,0 +1,60 @@ +""" +This module provides NeedleFiles class for interacting with Needle API's files endpoint. +""" + +import requests + +from needle.v1.models import ( + NeedleConfig, + NeedleBaseClient, + Error, +) + + +class NeedleFiles(NeedleBaseClient): + """ + A client for interacting with the Needle API's files endpoint. + + This class provides methods to create upload and download URLs for files within the Needle API. + It uses a requests session to handle HTTP requests with a default timeout of 120 seconds. + """ + + def __init__(self, config: NeedleConfig, headers: dict): + super().__init__(config, headers) + + self.endpoint = f"{config.url}/api/v1/files" + + # requests config + self.session = requests.Session() + self.session.headers.update(headers) + self.session.timeout = 120 + + def get_download_url(self, file_id: str) -> str: + """ + Retrieves the download URL for the given file. + If the file was manually uploaded, then the download URL will be valid for a short time therefore you should read the file before it expires. + + Args: + file_id (str): ID of the file to get the download URL for. + + Returns: + str: The download URL for the file. + + Raises: + Error: If the API request fails. + """ + if not file_id: + raise Error( + message="file_id is required", + code=422, + ) + + url = f"{self.endpoint}/{file_id}/download_url" + resp = self.session.get(url) + body = resp.json() + + if resp.status_code >= 400: + error = body.get("error") + raise Error(**error) + + return body.get("result") diff --git a/pyproject.toml b/pyproject.toml index a2886bd..997faf2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api" [tool.poetry] name = "needle-python" -version = "0.3.0" +version = "0.4.0" description = "Needle client library for Python" authors = [ "Onur Eken ", diff --git a/tutorials/needle.ipynb b/tutorials/collections.ipynb similarity index 79% rename from tutorials/needle.ipynb rename to tutorials/collections.ipynb index 6b20974..7eacbb0 100644 --- a/tutorials/needle.ipynb +++ b/tutorials/collections.ipynb @@ -35,13 +35,11 @@ "outputs": [], "source": [ "from needle.v1 import NeedleClient\n", - "from needle.v1.models import FileToAdd\n", - "\n", "\n", "ndl = NeedleClient()\n", - "collection_id = \"clt_01J4NW2THDSQC7GDD19CMCTAAK\"\n", + "collection_id = \"clt_01J6SPFD61D5QYGSHK6W15M3VB\"\n", "\n", - "prompt = \"What techniques moved into adopt in this volume of technology radar?\"\n", + "prompt = \"How do joins work in PQL?\"\n", "results = ndl.collections.search(collection_id, text=prompt)\n", "\n", "for r in results:\n", @@ -50,26 +48,18 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Retrieval-Augmented Generation (RAG) technique moved into the \"Adopt\" category in this volume of the Technology Radar.\n" - ] - } - ], + "outputs": [], "source": [ "system_messages = [{\"role\": \"system\", \"content\": r.content} for r in results] # results from Needle\n", "user_message = {\n", " \"role\": \"user\",\n", " \"content\": f\"\"\"\n", " Only answer the question based on the provided results data. \n", - " If there is no data in the provided data for the question, do not try to generate an answer. \n", + " If there is no data in the provided data for the question, do not generate an answer. \n", " This is the question: {prompt}\n", - "\"\"\",\n", + " \"\"\",\n", "}\n", "\n", "openai_client = OpenAI()\n", diff --git a/tutorials/files.ipynb b/tutorials/files.ipynb new file mode 100644 index 0000000..f32ec6a --- /dev/null +++ b/tutorials/files.ipynb @@ -0,0 +1,59 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import os" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "os.environ[\"NEEDLE_API_KEY\"] = \"\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from needle.v1 import NeedleClient\n", + "\n", + "ndl = NeedleClient(\n", + " url=\"http://localhost:3000\",\n", + ")\n", + "collection_id = \"clt_01JBS48E0M6YF2AN9ZAXVK8S2P\"\n", + "\n", + "ndl.files.get_download_url(\"fle_01JBS5EXXRGSQMCY7B6XMRXETE\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "needle-tutorial-Pi3Ihry5", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.4" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +}