revisited method parameters

Author: oeken

.gitignore

@@ -160,3 +160,5 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+.DS_Store

README.md

@@ -1 +1,97 @@
-# needle-python
+# Needle Python Library
+
+This Python library provides convenient acccess to Needle API. There are various methods and data types which, we believe will help you explore Needle API quickly. There may be some functionality available in REST API earlier than this Python library. In any case, we recommend to take look the the complete [documentation](https://docs.needle-ai.com). Thank you for flying with us. 🚀
+
+## Installation
+
+This library requires Python >3.8 and `pip` to use. You don't need the sources unless you want to modify it. Install with:
+
+```
+pip install needle-python
+```
+
+## Usage ⚡️
+
+To get started, generate an API key for your account in developer settings menu at (Needle)[https://needle-ai.com]. Note that your key will be valid until you revoke it. Set the following env variable before you run your code:
+
+```
+export NEEDLE_API_KEY=<your-api-key>
+```
+
+`NeedleClient` reads the API key from the environment by default. If you like to override this behaviour you can pass it in as a parameter. 
+
+### Retrieve context from Needle
+
+```python
+from needle.v1 import NeedleClient
+from needle.v1.models import FileToAdd
+
+
+ndl = NeedleClient()
+collection = ndl.collections.create(name="Tech Trends")
+
+# add file to collection
+files = ndl.collections.files.add(
+    collection_id=collection_id,
+    files=[
+        FileToAdd(
+            name="tech-radar-30.pdf",
+            url="https://www.thoughtworks.com/content/dam/thoughtworks/documents/radar/2024/04/tr_technology_radar_vol_30_en.pdf",
+        )
+    ],
+)
+
+# wait until indexing is complete
+files = ndl.collections.files.list(collection_id)
+if not all(f.status == "indexed" for f in files):
+    time.sleep(5)
+    files = ndl.collections.files.list(collection_id)
+
+# retrieve relevant context
+prompt = "What techniques moved into adopt in this volume of technology radar?"
+results = ndl.collections.search(collection_id, text=prompt)
+```
+
+Needle instantly extracts key points from your files.
+
+### Complete your RAG pipeline
+
+Naturally, to compose a human friendly answer use an LLM provider of your choice. For the demo purposes, we used OpenAI in this example:
+
+```python
+from openai import OpenAI
+
+system_messages = [{"role": "system", "content": r.content} for r in results] # results from Needle
+user_message = {
+    "role": "system",
+    "content": f"""
+    Do not hallucinate. Only answer the question based on the provided results data. 
+    If there is no data in the provided data for the question, do not try to generate an answer that does not make sense. 
+    This is the question: {prompt}
+""",
+}
+
+openai_client = OpenAI()
+answer = openai_client.chat.completions.create(
+    model="gpt-3.5-turbo",
+    messages=[
+        *system_messages,
+        user_message,
+    ],
+)
+
+print(answer.choices[0].message.content)
+# -> Retrieval-Augmented Generation (RAG) is the technique that moved into "Adopt" in this volume of the Technology Radar.
+```
+
+This is one basic example of a RAG pipeline you can quicklu implement using Needle and OpenAI. Feel free to engineer more precise prompts and explore other prompting techniques such as chain-of-thoughts (CoT), graph of thoughts (GoT) etc. 
+
+Needle API helps you with hassle-free contextualization however does not limit you to a certain RAG technique. Let us know what you build in our (Discord channel)[https://discord.gg/JzJcHgTyZx] :)
+
+## Exceptions 🧨
+
+If a request to Needle API fails, `needle.v1.models.Error` object will be thrown. There you can see a `message` and more details about the error.
+
+# Support 📞
+
+If you have questions you can contact us in our (Discord channel)[https://discord.gg/JzJcHgTyZx]. 

needle/__init__.py

@@ -0,0 +1,7 @@
+from urllib.parse import urlparse, urlunparse
+
+
+def make_needle_search_url(needle_url: str):
+    parsed_url = urlparse(needle_url)
+    new_netloc = f"search.{parsed_url.netloc}"
+    return urlunparse(parsed_url._replace(netloc=new_netloc))

needle/utils.py

@@ -0,0 +1,48 @@
+"""
+This module provides NeedlClient class for interacting with Needle API.
+"""
+
+from typing import Optional
+import os
+
+from needle.utils import make_needle_search_url
+from needle.v1.models import (
+    NeedleConfig,
+    NeedleBaseClient,
+)
+from needle.v1.collections import NeedleCollections
+
+
+NEEDLE_DEFAULT_URL = "https://needle-ai.com"
+
+
+class NeedleClient(NeedleBaseClient):
+    """
+    A client for interacting with the Needle API.
+
+    This class provides a high-level interface for interacting with the Needle API,
+    including managing collections and performing searches.
+
+    Initialize the client with an API key and an optional URL.
+    If no API key is provided, the client will use the `NEEDLE_API_KEY` environment variable.
+    If no URL is provided, the client will use the default Needle API URL, that is https://needle-ai.com.
+
+    Attributes:
+        collections (NeedleCollections): A client for managing collections within the Needle API.
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = os.environ.get("NEEDLE_API_KEY"),
+        url: Optional[str] = NEEDLE_DEFAULT_URL,
+        _search_url: Optional[str] = None,
+    ):
+        if not _search_url:
+            _search_url = make_needle_search_url(url)
+
+        config = NeedleConfig(api_key, url, search_url=_search_url)
+        headers = {"x-api-key": config.api_key}
+        super().__init__(config, headers)
+
+        # sub clients
+        self.collections = NeedleCollections(config, headers)

needle/v1/__init__.py

@@ -0,0 +1,125 @@
+"""
+This module provides NeedleCollections class for interacting with Needle API's collections endpoint.
+"""
+
+from typing import Optional
+
+import requests
+
+from needle.v1.models import (
+    NeedleConfig,
+    NeedleBaseClient,
+    Collection,
+    Error,
+    SearchResult,
+)
+from needle.v1.collections.files import NeedleCollectionsFiles
+
+
+class NeedleCollections(NeedleBaseClient):
+    """
+    A client for interacting with the Needle API's collections endpoint.
+
+    This class provides methods to create and manage collections 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/collections"
+        self.search_endpoint = f"{config.search_url}/api/v1/collections"
+
+        # requests config
+        self.session = requests.Session()
+        self.session.headers.update(headers)
+        self.session.timeout = 120
+
+        # sub clients
+        self.files = NeedleCollectionsFiles(config, headers)
+
+    def create(self, name: str, file_ids: Optional[list[str]] = None):
+        """
+        Creates a new collection with the specified name and file IDs.
+
+        Args:
+            name (str): The name of the collection.
+            file_ids (Optiona[list[str]]): A list of file IDs to include in the collection.
+
+        Returns:
+            Collection: The created collection object.
+
+        Raises:
+            Error: If the API request fails.
+        """
+        req_body = {"name": name, "file_ids": file_ids}
+        resp = self.session.post(
+            f"{self.endpoint}",
+            json=req_body,
+        )
+        body = resp.json()
+        if resp.status_code >= 400:
+            error = body.get("error")
+            raise Error(**error)
+        c = body.get("result")
+        return Collection(**c)
+
+    def get(self, collection_id: str):
+        """
+        Retrieves a collection by its ID.
+
+        Args:
+            collection_id (str): The ID of the collection to retrieve.
+
+        Returns:
+            Collection: The retrieved collection object.
+
+        Raises:
+            Error: If the API request fails.
+        """
+        resp = self.session.get(f"{self.endpoint}/{collection_id}")
+        body = resp.json()
+        if resp.status_code >= 400:
+            error = body.get("error")
+            raise Error(**error)
+        c = body.get("result")
+        return Collection(**c)
+
+    def list(self):
+        """
+        Lists all collections.
+
+        Returns:
+            list[Collection]: A list of all collections.
+
+        Raises:
+            Error: If the API request fails.
+        """
+        resp = self.session.get(self.endpoint)
+        body = resp.json()
+        if resp.status_code >= 400:
+            error = body.get("error")
+            raise Error(**error)
+        return [Collection(**c) for c in body.get("result")]
+
+    def search(self, collection_id: str, text: str):
+        """
+        Searches within a collection based on the provided parameters.
+
+        Args:
+            params (SearchCollectionRequest): The search parameters.
+
+        Returns:
+            list[dict]: The search results.
+
+        Raises:
+            Error: If the API request fails.
+        """
+        endpoint = f"{self.search_endpoint}/{collection_id}/search"
+        req_body = {"text": text}
+        resp = self.session.post(endpoint, headers=self.headers, json=req_body)
+        body = resp.json()
+        if resp.status_code >= 400:
+            error = body.get("error")
+            raise Error(**error)
+        return [SearchResult(**c) for c in body.get("result")]

needle/v1/collections/__init__.py

@@ -0,0 +1,74 @@
+"""
+This module provides NeedleCollectionsFiles class for interacting with 
+Needle API's collectiton files endpoint.
+"""
+
+from dataclasses import asdict
+import requests
+
+from needle.v1 import NeedleBaseClient, NeedleConfig
+from needle.v1.models import FileToAdd, Error, CollectionFile
+
+
+class NeedleCollectionsFiles(NeedleBaseClient):
+    """
+    A client for interacting with the Needle API's collection files endpoint.
+
+    This class provides methods to create and manage collection 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.collections_endpoint = f"{config.url}/api/v1/collections"
+
+        # requests config
+        self.session = requests.Session()
+        self.session.headers.update(headers)
+        self.session.timeout = 120
+
+    def add(self, collection_id: str, files: list[FileToAdd]):
+        """
+        Adds files to a specified collection. Added files will be automatically indexed and after be available for search within the collection.
+
+        Args:
+            collection_id (str): The ID of the collection to which files will be added.
+            files (list[FileToAdd]): A list of FileToAdd objects representing the files to be added.
+
+        Returns:
+            list[CollectionFile]: A list of CollectionFile objects representing the added files.
+
+        Raises:
+            Error: If the API request fails.
+        """
+
+        endpoint = f"{self.collections_endpoint}/{collection_id}/files"
+        req_body = {"files": [asdict(f) for f in files]}
+        resp = self.session.post(endpoint, json=req_body)
+        body = resp.json()
+        if resp.status_code >= 400:
+            error = body.get("error")
+            raise Error(**error)
+        return [CollectionFile(**cf) for cf in body.get("result")]
+
+    def list(self, collection_id: str):
+        """
+        Lists all files in a specified collection.
+
+        Args:
+            collection_id (str): The ID of the collection whose files will be listed.
+
+        Returns:
+            list[CollectionFile]: A list of CollectionFile objects representing the files in the collection.
+
+        Raises:
+            Error: If the API request fails.
+        """
+        endpoint = f"{self.collections_endpoint}/{collection_id}/files"
+        resp = self.session.get(endpoint)
+        body = resp.json()
+        if resp.status_code >= 400:
+            error = body.get("error")
+            raise Error(**error)
+        return [CollectionFile(**cf) for cf in body.get("result")]