Add NeedleFiles sub-client (#3)

Author: oeken

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)

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")

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 <m.onureken@gmail.com>",

tutorials/needle.ipynb tutorials/collections.ipynbtutorials/needle.ipynb renamed to 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",

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
+}