needlesearchai 2.1.0__py3-none-any.whl

needlesearchai/__init__.py

@@ -0,0 +1,78 @@
+"""NeedleSearch Python SDK.
+
+Quick start:
+    from needlesearchai import NeedleSearch
+
+    ns = NeedleSearch(api_key="nsk_...")
+    results = ns.search("termination clauses")                      # vector search
+    overview = ns.search("notice period?", include_overview=True)   # search + AI answer
+    for chunk in ns.chat.stream("Explain promissory estoppel"):     # plain chat (no docs)
+        print(chunk.text, end="", flush=True)
+
+    for item in ns.items.iterate():                                 # browse, cursor-paged
+        print(item.name, item.status)
+
+    job = ns.jobs.create("Summarise the risks", item_ids=[doc_id])  # async research
+    print(ns.jobs.wait(job.id).answer)
+"""
+from needlesearchai.client import NeedleSearch
+from needlesearchai.exceptions import (
+    AuthError,
+    DocumentNotReadyError,
+    InsufficientScopeError,
+    NeedleSearchError,
+    NotFoundError,
+    QuotaError,
+    RateLimitError,
+    ServerError,
+)
+from needlesearchai.types import (
+    ChatChunk,
+    FileRecord,
+    Item,
+    ItemPage,
+    MeInfo,
+    QuotaInfo,
+    ResearchJob,
+    ResearchResult,
+    SearchResponse,
+    SearchResult,
+    ServiceStatus,
+    UploadJob,
+    UploadJobPage,
+    UploadResult,
+    Webhook,
+    WebhookDelivery,
+)
+from needlesearchai.webhooks import verify_signature
+
+__all__ = [
+    "NeedleSearch",
+    "NeedleSearchError",
+    "AuthError",
+    "InsufficientScopeError",
+    "RateLimitError",
+    "QuotaError",
+    "NotFoundError",
+    "ServerError",
+    "DocumentNotReadyError",
+    "SearchResult",
+    "SearchResponse",
+    "ChatChunk",
+    "FileRecord",
+    "UploadResult",
+    "QuotaInfo",
+    "MeInfo",
+    "ResearchResult",
+    "Item",
+    "ItemPage",
+    "ResearchJob",
+    "ServiceStatus",
+    "UploadJob",
+    "UploadJobPage",
+    "Webhook",
+    "WebhookDelivery",
+    "verify_signature",
+]
+
+__version__ = "2.1.0"

needlesearchai/_http.py

@@ -0,0 +1,138 @@
+"""HTTP client wrapper with auth, retry, and error handling."""
+from __future__ import annotations
+
+import time
+from typing import Any, Generator
+
+try:
+    import httpx
+except ImportError as e:
+    raise ImportError("needlesearchai requires httpx. Install with: pip install needlesearchai[http]") from e
+
+from needlesearchai.exceptions import (
+    AuthError,
+    DocumentNotReadyError,
+    InsufficientScopeError,
+    NeedleSearchError,
+    NotFoundError,
+    QuotaError,
+    RateLimitError,
+    ServerError,
+)
+
+_DEFAULT_TIMEOUT = 30.0
+_STREAM_TIMEOUT  = 300.0
+_MAX_RETRIES     = 3
+
+
+def _raise_for_status(response: httpx.Response) -> None:
+    if response.is_success:
+        return
+
+    try:
+        body = response.json()
+    except Exception:
+        body = {}
+
+    error_code = body.get("error", "")
+    try:
+        raw_text = response.text[:300]
+    except Exception:
+        raw_text = ""
+    message = body.get("message", body.get("detail", raw_text))
+    status     = response.status_code
+
+    if status == 401:
+        raise AuthError(message or "Invalid or missing API key", status_code=status, error_code=error_code)
+    if status == 403:
+        if error_code == "insufficient_scope":
+            raise InsufficientScopeError(message, status_code=status, error_code=error_code)
+        raise AuthError(message, status_code=status, error_code=error_code)
+    if status == 404:
+        raise NotFoundError(message, status_code=status, error_code=error_code)
+    if status == 429:
+        if error_code == "monthly_quota_exceeded":
+            raise QuotaError(message, used=body.get("used"), limit=body.get("limit"))
+        retry_after = int(response.headers.get("Retry-After", 60))
+        raise RateLimitError(message, retry_after=retry_after)
+    if status == 400:
+        if error_code == "document_not_ready":
+            raise DocumentNotReadyError(message, status_code=status, error_code=error_code)
+        raise NeedleSearchError(message, status_code=status, error_code=error_code)
+    if status >= 500:
+        raise ServerError(message or f"Server error {status}", status_code=status)
+
+    raise NeedleSearchError(message or f"HTTP {status}", status_code=status, error_code=error_code)
+
+
+class HttpClient:
+    """Synchronous HTTP client with exponential backoff retry."""
+
+    def __init__(self, api_key: str, base_url: str, timeout: float = _DEFAULT_TIMEOUT):
+        self._client = httpx.Client(
+            base_url=base_url,
+            headers={"Authorization": f"Bearer {api_key}", "User-Agent": "needlesearchai-python/1.0"},
+            timeout=timeout,
+        )
+
+    def get(self, path: str, **params) -> Any:
+        return self._request("GET", path, params=params or None)
+
+    def post(self, path: str, json: Any = None, *, idempotency_key: str | None = None) -> Any:
+        headers = {"Idempotency-Key": idempotency_key} if idempotency_key else None
+        return self._request("POST", path, json=json, headers=headers)
+
+    def patch(self, path: str, json: Any = None) -> Any:
+        return self._request("PATCH", path, json=json)
+
+    def delete(self, path: str) -> Any:
+        return self._request("DELETE", path)
+
+    def _request(self, method: str, path: str, **kwargs) -> Any:
+        for attempt in range(_MAX_RETRIES):
+            try:
+                # Drop unset optionals so httpx does not send empty headers.
+                call = {k: v for k, v in kwargs.items() if v is not None}
+                r = self._client.request(method, path, **call)
+                _raise_for_status(r)
+                return r.json() if r.content else None
+            except RateLimitError as e:
+                if attempt == _MAX_RETRIES - 1:
+                    raise
+                wait = min(e.retry_after, 30)
+                time.sleep(wait)
+            except (httpx.TimeoutException, httpx.NetworkError) as e:
+                if attempt == _MAX_RETRIES - 1:
+                    raise NeedleSearchError(f"Network error: {e}") from e
+                time.sleep(2 ** attempt)
+
+    def raw(self, method: str, path: str, **kwargs):
+        """A checked response with its body intact — for endpoints that return
+        bytes rather than JSON (file downloads)."""
+        r = self._client.request(method, path, **kwargs)
+        _raise_for_status(r)
+        return r
+
+    def stream(self, path: str, json: Any = None) -> Generator[str, None, None]:
+        """Yield raw SSE data lines from a streaming endpoint."""
+        with self._client.stream("POST", path, json=json, timeout=_STREAM_TIMEOUT) as r:
+            _raise_for_status(r)
+            parts: list[str] = []
+            for line in r.iter_lines():
+                if line.startswith("data:"):
+                    raw = line[5:]
+                    parts.append(raw[1:] if raw.startswith(" ") else raw)
+                elif line == "" and parts:
+                    yield "\n".join(parts)
+                    parts = []
+            if parts:
+                yield "\n".join(parts)
+
+    def close(self):
+        self._client.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()

needlesearchai/chat.py

@@ -0,0 +1,69 @@
+"""NeedleSearch SDK — Chat API with streaming support."""
+from __future__ import annotations
+import json
+from typing import Generator
+
+from needlesearchai.types import ChatChunk
+
+
+class ChatAPI:
+    """Wrapper for POST /v1/chat."""
+
+    def __init__(self, http):
+        self._http = http
+
+    def stream(
+        self,
+        message: str,
+        *,
+        provider: str | None = None,
+    ) -> Generator[ChatChunk, None, None]:
+        """Stream a plain chat response from the AI.
+
+        Conversational chat only — it does NOT read your documents. For
+        document-grounded answers use ``client.search(..., include_overview=True)``
+        or the research API.
+
+        Args:
+            message:   User question or instruction
+            provider:  LLM provider override ('deepseek', 'openai', 'anthropic')
+
+        Yields:
+            ChatChunk with .text containing the streaming text
+
+        Example:
+            for chunk in ns.chat.stream("Explain promissory estoppel"):
+                print(chunk.text, end="", flush=True)
+        """
+        body = {"message": message}
+        if provider:
+            body["provider"] = provider
+
+        for data_line in self._http.stream("/chat", json=body):
+            if not data_line:
+                continue
+            try:
+                payload = json.loads(data_line)
+                # Only yield actual text delta events
+                if isinstance(payload, str):
+                    yield ChatChunk(text=payload)
+                elif isinstance(payload, dict) and "text" in payload:
+                    yield ChatChunk(text=payload["text"])
+                elif isinstance(payload, dict) and payload.get("event") == "done":
+                    break
+            except (json.JSONDecodeError, TypeError):
+                # Plain text delta
+                if data_line and not data_line.startswith("{"):
+                    yield ChatChunk(text=data_line)
+
+    def complete(self, message: str, **kwargs) -> str:
+        """Non-streaming chat — collects all chunks and returns the full text.
+
+        Args:
+            message: User question or instruction
+            **kwargs: Same keyword arguments as stream()
+
+        Returns:
+            Complete AI response as a string
+        """
+        return "".join(chunk.text for chunk in self.stream(message, **kwargs))

needlesearchai/client.py

@@ -0,0 +1,114 @@
+"""NeedleSearch Python SDK — main client."""
+from __future__ import annotations
+
+from needlesearchai._http import HttpClient
+from needlesearchai.chat import ChatAPI
+from needlesearchai.files import FilesAPI
+from needlesearchai.folders import FoldersAPI
+from needlesearchai.items import ItemsAPI
+from needlesearchai.jobs import ResearchJobsAPI
+from needlesearchai.research import ResearchAPI
+from needlesearchai.search import SearchAPI
+from needlesearchai.types import MeInfo, QuotaInfo, ServiceStatus
+from needlesearchai.uploads import UploadsAPI
+from needlesearchai.webhooks import WebhooksAPI
+
+_DEFAULT_BASE_URL = "https://your-instance.com/v1"
+
+
+class NeedleSearch:
+    """NeedleSearch API client.
+
+    Usage:
+        from needlesearchai import NeedleSearch
+
+        ns = NeedleSearch(api_key="nsk_...")
+
+        # Semantic search (optionally with an AI overview over your documents)
+        results = ns.search("What are the termination clauses?")
+        overview = ns.search("How long is the notice period?", include_overview=True)
+
+        # Streaming chat — plain conversation, does NOT read your documents
+        for chunk in ns.chat.stream("Explain promissory estoppel"):
+            print(chunk.text, end="", flush=True)
+
+        # Deep research over your documents (Starter+ plan)
+        answer = ns.research.complete("What are the risks in this contract?", document_ids=["..."])
+    """
+
+    def __init__(self, api_key: str, base_url: str = _DEFAULT_BASE_URL, timeout: float = 30.0):
+        """Create a NeedleSearch client.
+
+        Args:
+            api_key:  Your API key (starts with 'nsk_')
+            base_url: API base URL (default: https://your-instance.com/v1)
+            timeout:  Request timeout in seconds
+        """
+        if not api_key:
+            raise ValueError("api_key is required")
+        if not api_key.startswith("nsk_"):
+            raise ValueError("api_key must start with 'nsk_'")
+
+        self._http = HttpClient(api_key=api_key, base_url=base_url.rstrip("/"), timeout=timeout)
+
+        self.search   = SearchAPI(self._http)
+        self.chat     = ChatAPI(self._http)
+        self.research = ResearchAPI(self._http)
+        self.files    = FilesAPI(self._http)
+        self.items    = ItemsAPI(self._http)
+        self.folders  = FoldersAPI(self._http)
+        self.webhooks = WebhooksAPI(self._http)
+        self.jobs     = ResearchJobsAPI(self._http)
+        self.uploads  = UploadsAPI(self._http)
+
+    def status(self) -> ServiceStatus:
+        """Whether the API can currently serve each class of operation.
+
+        The one call that needs no key — so it still answers when the key itself
+        is the suspect. `down` means retry with backoff rather than debugging
+        your integration.
+        """
+        data = self._http.get("/status")
+        return ServiceStatus(
+            status=data.get("status", "unknown"),
+            services=data.get("services", {}),
+            api_version=data.get("api_version", "v1"),
+            checked_at=data.get("checked_at"),
+        )
+
+    def me(self) -> MeInfo:
+        """Return identity and plan information for this API key."""
+        data = self._http.get("/me")
+        return MeInfo(
+            key_prefix=data.get("key_prefix", ""),
+            name=data.get("name", ""),
+            plan=data.get("plan", "free"),
+            scopes=data.get("scopes", []),
+            organization_id=data.get("organization_id", ""),
+            features=data.get("features", {}),
+            created_at=data.get("created_at"),
+        )
+
+    def quota(self) -> QuotaInfo:
+        """Return current quota and plan status."""
+        data = self._http.get("/quota")
+        monthly = data.get("monthly_requests", {})
+        return QuotaInfo(
+            plan=data.get("plan", "free"),
+            monthly_limit=monthly.get("limit"),
+            monthly_used=monthly.get("used", 0),
+            monthly_remaining=monthly.get("remaining"),
+            rate_limit_rpm=data.get("rate_limit_rpm", 0),
+            features=data.get("features", {}),
+            reset_at=data.get("reset_at"),
+        )
+
+    def close(self) -> None:
+        """Close the underlying HTTP client."""
+        self._http.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()

needlesearchai/exceptions.py

@@ -0,0 +1,47 @@
+"""NeedleSearch SDK exceptions."""
+
+
+class NeedleSearchError(Exception):
+    """Base exception for all NeedleSearch errors."""
+    def __init__(self, message: str, status_code: int | None = None, error_code: str | None = None):
+        super().__init__(message)
+        self.status_code = status_code
+        self.error_code = error_code
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.args[0]!r}, status_code={self.status_code})"
+
+
+class AuthError(NeedleSearchError):
+    """API key missing, invalid, revoked, or expired."""
+
+
+class InsufficientScopeError(NeedleSearchError):
+    """API key lacks the required scope for this operation."""
+
+
+class RateLimitError(NeedleSearchError):
+    """Rate limit exceeded. Check retry_after for seconds to wait."""
+    def __init__(self, message: str, retry_after: int = 60):
+        super().__init__(message, status_code=429, error_code="rate_limit_exceeded")
+        self.retry_after = retry_after
+
+
+class QuotaError(NeedleSearchError):
+    """Monthly request quota exceeded."""
+    def __init__(self, message: str, used: int | None = None, limit: int | None = None):
+        super().__init__(message, status_code=429, error_code="monthly_quota_exceeded")
+        self.used = used
+        self.limit = limit
+
+
+class NotFoundError(NeedleSearchError):
+    """Requested resource not found."""
+
+
+class ServerError(NeedleSearchError):
+    """Internal server error."""
+
+
+class DocumentNotReadyError(NeedleSearchError):
+    """Document is still being processed."""

needlesearchai/files.py

@@ -0,0 +1,187 @@
+"""NeedleSearch SDK — Files API."""
+from __future__ import annotations
+from pathlib import Path
+from typing import BinaryIO
+
+from needlesearchai.exceptions import NeedleSearchError
+from needlesearchai.types import FileRecord, UploadResult
+
+
+class FilesAPI:
+    """Wrapper for /v1/files endpoints.
+
+    Requires Pro plan or higher.
+    """
+
+    def __init__(self, http):
+        self._http = http
+
+    def list(self, *, limit: int = 50, offset: int = 0) -> list[FileRecord]:
+        """Files at the workspace root.
+
+        Deprecated in favour of `client.items.list()`, which returns folders too
+        and pages by cursor. This called `GET /v1/files`, an endpoint that no
+        longer exists — every call raised NotFoundError — so it is routed
+        through /v1/items and filtered to files.
+        """
+        data = self._http.get("/items", limit=limit, offset=offset)
+        return [
+            _parse_file(i) for i in data.get("items", []) if i.get("type") == "file"
+        ]
+
+    def get(self, file_id: str) -> FileRecord:
+        """Metadata for one file.
+
+        Deprecated in favour of `client.items.get()`. Also pointed at a removed
+        endpoint (`GET /v1/files/{id}`); it uses /v1/items/{id} now.
+        """
+        return _parse_file(self._http.get(f"/items/{file_id}"))
+
+    def delete(self, *file_ids: str, force: bool = False) -> dict:
+        """Delete files and/or folders by id — the type is detected per id.
+
+        Returns `{"deleted": [...], "failed": [...]}`. A document published in a
+        market dataset comes back under `failed` unless `force=True`.
+        """
+        if not file_ids:
+            raise ValueError("at least one id is required")
+        return self._http.post(
+            "/files/delete", json={"ids": list(file_ids), "force": force}
+        )
+
+    def reprocess(self, *file_ids: str) -> dict:
+        """Retry failed documents in place — no re-upload, no duplicate.
+
+        Ids that are not in a failed state come back under `failed`.
+        """
+        if not file_ids:
+            raise ValueError("at least one id is required")
+        return self._http.post("/files/reprocess", json={"ids": list(file_ids)})
+
+    def download(self, file_id: str) -> bytes:
+        """The original bytes. The document must have finished processing."""
+        response = self._http.raw("GET", f"/files/{file_id}/download")
+        return response.content
+
+    def upload(
+        self,
+        file: Path | BinaryIO | bytes,
+        *,
+        filename: str | None = None,
+        folder_id: str | None = None,
+        timeout: float = 120.0,
+    ) -> UploadResult:
+        """Upload a document for async processing.
+
+        Uses the resumable tus protocol under the hood: it creates an upload job
+        (`POST /v1/uploads`), streams the bytes to the tus endpoint, seals the
+        job, then polls until the document row exists. ZIP/archive files are
+        expanded server-side into a folder of documents.
+
+        Args:
+            file:      File path, file-like object, or bytes
+            filename:  Override filename (required if file is bytes or BinaryIO)
+            folder_id: Place document in this folder
+            timeout:   Seconds to wait for the byte upload to finalize
+
+        Returns:
+            UploadResult with the new document ID and status
+        """
+        import base64
+        import hashlib
+        import time
+        import uuid
+
+        try:
+            import httpx as _httpx
+        except ImportError:
+            raise ImportError("needlesearchai[http] required for file uploads")
+
+        if isinstance(file, Path):
+            filename = filename or file.name
+            content = file.read_bytes()
+        elif isinstance(file, bytes):
+            content = file
+        else:
+            content = file.read()
+
+        if not filename:
+            raise ValueError("filename is required when uploading bytes or a file-like object")
+
+        client_file_id = uuid.uuid4().hex
+        digest = hashlib.sha256(content).hexdigest()
+
+        # 1. Create the job + a single file slot.
+        job = self._http.post("/uploads", json={
+            "client_request_id": uuid.uuid4().hex,
+            "files": [{
+                "client_file_id": client_file_id,
+                "name": filename,
+                "size": len(content),
+                "hash": digest,
+            }],
+            **({"folder_id": folder_id} if folder_id else {}),
+        })
+        slot = next((f for f in job["files"] if f["client_file_id"] == client_file_id), job["files"][0])
+        token = slot["upload_token"]
+        tus_endpoint = job["tus_endpoint"]
+
+        # 2. tus upload — the token in Upload-Metadata is the auth (no Bearer on
+        #    the tus host). One create + one PATCH; resumable by design.
+        def _b64(s: str) -> str:
+            return base64.b64encode(s.encode()).decode()
+
+        meta = f"upload_token {_b64(token)},filename {_b64(filename)}"
+        with _httpx.Client(timeout=timeout) as tus:
+            cr = tus.post(tus_endpoint, headers={
+                "Tus-Resumable": "1.0.0",
+                "Upload-Length": str(len(content)),
+                "Upload-Metadata": meta,
+            })
+            cr.raise_for_status()
+            location = cr.headers.get("Location") or ""
+            if not location:
+                raise NeedleSearchError("tus create did not return a Location header")
+            if location.startswith("/"):
+                base = tus_endpoint.split("/", 3)
+                location = f"{base[0]}//{base[2]}{location}"
+            pr = tus.patch(location, content=content, headers={
+                "Tus-Resumable": "1.0.0",
+                "Upload-Offset": "0",
+                "Content-Type": "application/offset+octet-stream",
+            })
+            pr.raise_for_status()
+
+        # 3. Seal the job.
+        self._http.post(f"/uploads/{job['id']}/seal")
+
+        # 4. Poll until the document row exists (post-finish finalized the bytes).
+        deadline = time.time() + timeout
+        terminal = {"uploaded", "processing", "ready", "duplicate", "failed"}
+        last = slot
+        while time.time() < deadline:
+            prog = self._http.get(f"/uploads/{job['id']}")
+            last = next((f for f in prog.get("files", []) if f["client_file_id"] == client_file_id), last)
+            if last.get("document_id") or last.get("status") in terminal:
+                break
+            time.sleep(1.0)
+
+        if last.get("status") == "failed":
+            raise NeedleSearchError(f"upload failed for {filename}")
+        return UploadResult(
+            id=last.get("document_id") or "",
+            name=filename,
+            status=last.get("status", "pending"),
+            message="" if last.get("document_id") else "still finalizing; poll items API",
+        )
+
+def _parse_file(d: dict) -> FileRecord:
+    return FileRecord(
+        id=d.get("id", ""),
+        name=d.get("name", ""),
+        status=d.get("status", ""),
+        size_bytes=d.get("size_bytes"),
+        folder_id=d.get("folder_id"),
+        created_at=d.get("created_at"),
+        updated_at=d.get("updated_at"),
+    )

needlesearchai/folders.py

@@ -0,0 +1,30 @@
+"""NeedleSearch SDK — Folders API (`/v1/folders`)."""
+from __future__ import annotations
+
+from typing import Optional
+
+from needlesearchai.types import Item
+
+
+class FoldersAPI:
+    """Wrapper for /v1/folders. Requires the `files:write` scope."""
+
+    def __init__(self, http):
+        self._http = http
+
+    def create(self, name: str, *, parent_id: Optional[str] = None) -> Item:
+        """Create a folder, or return the existing one with the same name in the
+        same place. Idempotent server-side, so a retry is safe without a key.
+        """
+        body: dict = {"name": name}
+        if parent_id is not None:
+            body["parent_id"] = parent_id
+        data = self._http.post("/folders", json=body)
+        return Item(
+            id=data.get("id", ""),
+            type="folder",
+            name=data.get("name", ""),
+            parent_id=data.get("parent_id"),
+            item_count=data.get("item_count"),
+            created_at=data.get("created_at"),
+        )