ragit 0.7.2__py3-none-any.whl → 0.7.4__py3-none-any.whl

ragit/providers/ollama.py

@@ -7,9 +7,19 @@ Ollama provider for LLM and Embedding operations.
 
 This provider connects to a local or remote Ollama server.
 Configuration is loaded from environment variables.
+
+Performance optimizations:
+- Connection pooling via requests.Session()
+- Async parallel embedding via trio + httpx
+- LRU cache for repeated embedding queries
 """
 
+from functools import lru_cache
+from typing import Any
+
+import httpx
 import requests
+import trio
 
 from ragit.config import config
 from ragit.providers.base import (
@@ -20,10 +30,37 @@ from ragit.providers.base import (
 )
 
 
+# Module-level cache for embeddings (shared across instances)
+@lru_cache(maxsize=2048)
+def _cached_embedding(text: str, model: str, embedding_url: str, timeout: int) -> tuple[float, ...]:
+    """Cache embedding results to avoid redundant API calls."""
+    # Truncate oversized inputs
+    if len(text) > OllamaProvider.MAX_EMBED_CHARS:
+        text = text[: OllamaProvider.MAX_EMBED_CHARS]
+
+    response = requests.post(
+        f"{embedding_url}/api/embeddings",
+        headers={"Content-Type": "application/json"},
+        json={"model": model, "prompt": text},
+        timeout=timeout,
+    )
+    response.raise_for_status()
+    data = response.json()
+    embedding = data.get("embedding", [])
+    if not embedding:
+        raise ValueError("Empty embedding returned from Ollama")
+    return tuple(embedding)
+
+
 class OllamaProvider(BaseLLMProvider, BaseEmbeddingProvider):
     """
     Ollama provider for both LLM and Embedding operations.
 
+    Performance features:
+    - Connection pooling via requests.Session() for faster sequential requests
+    - Async parallel embedding via embed_batch_async() using trio + httpx
+    - LRU cache for repeated embedding queries (2048 entries)
+
     Parameters
     ----------
     base_url : str, optional
@@ -32,6 +69,8 @@ class OllamaProvider(BaseLLMProvider, BaseEmbeddingProvider):
         API key for authentication (default: from OLLAMA_API_KEY env var)
     timeout : int, optional
         Request timeout in seconds (default: from OLLAMA_TIMEOUT env var)
+    use_cache : bool, optional
+        Enable embedding cache (default: True)
 
     Examples
     --------
@@ -39,12 +78,12 @@ class OllamaProvider(BaseLLMProvider, BaseEmbeddingProvider):
     >>> response = provider.generate("What is RAG?", model="llama3")
     >>> print(response.text)
 
-    >>> embedding = provider.embed("Hello world", model="nomic-embed-text")
-    >>> print(len(embedding.embedding))
+    >>> # Async batch embedding (5-10x faster for large batches)
+    >>> embeddings = trio.run(provider.embed_batch_async, texts, "mxbai-embed-large")
     """
 
     # Known embedding model dimensions
-    EMBEDDING_DIMENSIONS = {
+    EMBEDDING_DIMENSIONS: dict[str, int] = {
         "nomic-embed-text": 768,
         "nomic-embed-text:latest": 768,
         "mxbai-embed-large": 1024,
@@ -56,20 +95,48 @@ class OllamaProvider(BaseLLMProvider, BaseEmbeddingProvider):
         "qwen3-embedding:8b": 4096,
     }
 
+    # Max characters per embedding request (safe limit for 512 token models)
+    MAX_EMBED_CHARS = 2000
+
     def __init__(
         self,
         base_url: str | None = None,
         embedding_url: str | None = None,
         api_key: str | None = None,
         timeout: int | None = None,
+        use_cache: bool = True,
     ) -> None:
         self.base_url = (base_url or config.OLLAMA_BASE_URL).rstrip("/")
         self.embedding_url = (embedding_url or config.OLLAMA_EMBEDDING_URL).rstrip("/")
         self.api_key = api_key or config.OLLAMA_API_KEY
         self.timeout = timeout or config.OLLAMA_TIMEOUT
+        self.use_cache = use_cache
         self._current_embed_model: str | None = None
         self._current_dimensions: int = 768  # default
 
+        # Connection pooling via session
+        self._session: requests.Session | None = None
+
+    @property
+    def session(self) -> requests.Session:
+        """Lazy-initialized session for connection pooling."""
+        if self._session is None:
+            self._session = requests.Session()
+            self._session.headers.update({"Content-Type": "application/json"})
+            if self.api_key:
+                self._session.headers.update({"Authorization": f"Bearer {self.api_key}"})
+        return self._session
+
+    def close(self) -> None:
+        """Close the session and release resources."""
+        if self._session is not None:
+            self._session.close()
+            self._session = None
+
+    def __del__(self) -> None:
+        """Cleanup on garbage collection."""
+        self.close()
+
     def _get_headers(self, include_auth: bool = True) -> dict[str, str]:
         """Get request headers including authentication if API key is set."""
         headers = {"Content-Type": "application/json"}
@@ -88,21 +155,19 @@ class OllamaProvider(BaseLLMProvider, BaseEmbeddingProvider):
     def is_available(self) -> bool:
         """Check if Ollama server is reachable."""
         try:
-            response = requests.get(
+            response = self.session.get(
                 f"{self.base_url}/api/tags",
-                headers=self._get_headers(),
                 timeout=5,
             )
             return response.status_code == 200
         except requests.RequestException:
             return False
 
-    def list_models(self) -> list[dict[str, str]]:
+    def list_models(self) -> list[dict[str, Any]]:
         """List available models on the Ollama server."""
         try:
-            response = requests.get(
+            response = self.session.get(
                 f"{self.base_url}/api/tags",
-                headers=self._get_headers(),
                 timeout=10,
             )
             response.raise_for_status()
@@ -135,9 +200,8 @@ class OllamaProvider(BaseLLMProvider, BaseEmbeddingProvider):
             payload["system"] = system_prompt
 
         try:
-            response = requests.post(
+            response = self.session.post(
                 f"{self.base_url}/api/generate",
-                headers=self._get_headers(),
                 json=payload,
                 timeout=self.timeout,
             )
@@ -158,29 +222,34 @@ class OllamaProvider(BaseLLMProvider, BaseEmbeddingProvider):
             raise ConnectionError(f"Ollama generate failed: {e}") from e
 
     def embed(self, text: str, model: str) -> EmbeddingResponse:
-        """Generate embedding using Ollama (uses embedding_url, no auth for local)."""
+        """Generate embedding using Ollama with optional caching."""
         self._current_embed_model = model
         self._current_dimensions = self.EMBEDDING_DIMENSIONS.get(model, 768)
 
         try:
-            response = requests.post(
-                f"{self.embedding_url}/api/embeddings",
-                headers=self._get_headers(include_auth=False),
-                json={"model": model, "prompt": text},
-                timeout=self.timeout,
-            )
-            response.raise_for_status()
-            data = response.json()
-
-            embedding = data.get("embedding", [])
-            if not embedding:
-                raise ValueError("Empty embedding returned from Ollama")
+            if self.use_cache:
+                # Use cached version
+                embedding = _cached_embedding(text, model, self.embedding_url, self.timeout)
+            else:
+                # Direct call without cache
+                truncated = text[: self.MAX_EMBED_CHARS] if len(text) > self.MAX_EMBED_CHARS else text
+                response = self.session.post(
+                    f"{self.embedding_url}/api/embeddings",
+                    json={"model": model, "prompt": truncated},
+                    timeout=self.timeout,
+                )
+                response.raise_for_status()
+                data = response.json()
+                embedding_list = data.get("embedding", [])
+                if not embedding_list:
+                    raise ValueError("Empty embedding returned from Ollama")
+                embedding = tuple(embedding_list)
 
             # Update dimensions from actual response
             self._current_dimensions = len(embedding)
 
             return EmbeddingResponse(
-                embedding=tuple(embedding),
+                embedding=embedding,
                 model=model,
                 provider=self.provider_name,
                 dimensions=len(embedding),
@@ -189,7 +258,9 @@ class OllamaProvider(BaseLLMProvider, BaseEmbeddingProvider):
             raise ConnectionError(f"Ollama embed failed: {e}") from e
 
     def embed_batch(self, texts: list[str], model: str) -> list[EmbeddingResponse]:
-        """Generate embeddings for multiple texts (uses embedding_url, no auth for local).
+        """Generate embeddings for multiple texts sequentially.
+
+        For better performance with large batches, use embed_batch_async().
 
         Note: Ollama /api/embeddings only supports single prompts, so we loop.
         """
@@ -199,22 +270,28 @@ class OllamaProvider(BaseLLMProvider, BaseEmbeddingProvider):
         results = []
         try:
             for text in texts:
-                response = requests.post(
-                    f"{self.embedding_url}/api/embeddings",
-                    headers=self._get_headers(include_auth=False),
-                    json={"model": model, "prompt": text},
-                    timeout=self.timeout,
-                )
-                response.raise_for_status()
-                data = response.json()
+                # Truncate oversized inputs
+                truncated = text[: self.MAX_EMBED_CHARS] if len(text) > self.MAX_EMBED_CHARS else text
+
+                if self.use_cache:
+                    embedding = _cached_embedding(truncated, model, self.embedding_url, self.timeout)
+                else:
+                    response = self.session.post(
+                        f"{self.embedding_url}/api/embeddings",
+                        json={"model": model, "prompt": truncated},
+                        timeout=self.timeout,
+                    )
+                    response.raise_for_status()
+                    data = response.json()
+                    embedding_list = data.get("embedding", [])
+                    embedding = tuple(embedding_list) if embedding_list else ()
 
-                embedding = data.get("embedding", [])
                 if embedding:
                     self._current_dimensions = len(embedding)
 
                 results.append(
                     EmbeddingResponse(
-                        embedding=tuple(embedding),
+                        embedding=embedding,
                         model=model,
                         provider=self.provider_name,
                         dimensions=len(embedding),
@@ -224,6 +301,87 @@ class OllamaProvider(BaseLLMProvider, BaseEmbeddingProvider):
         except requests.RequestException as e:
             raise ConnectionError(f"Ollama batch embed failed: {e}") from e
 
+    async def embed_batch_async(
+        self,
+        texts: list[str],
+        model: str,
+        max_concurrent: int = 10,
+    ) -> list[EmbeddingResponse]:
+        """Generate embeddings for multiple texts in parallel using trio.
+
+        This method is 5-10x faster than embed_batch() for large batches
+        by making concurrent HTTP requests.
+
+        Parameters
+        ----------
+        texts : list[str]
+            Texts to embed.
+        model : str
+            Embedding model name.
+        max_concurrent : int
+            Maximum concurrent requests (default: 10).
+            Higher values = faster but more server load.
+
+        Returns
+        -------
+        list[EmbeddingResponse]
+            Embeddings in the same order as input texts.
+
+        Examples
+        --------
+        >>> import trio
+        >>> embeddings = trio.run(provider.embed_batch_async, texts, "mxbai-embed-large")
+        """
+        self._current_embed_model = model
+        self._current_dimensions = self.EMBEDDING_DIMENSIONS.get(model, 768)
+
+        # Results storage (index -> embedding)
+        results: dict[int, EmbeddingResponse] = {}
+        errors: list[Exception] = []
+
+        # Semaphore to limit concurrency
+        limiter = trio.CapacityLimiter(max_concurrent)
+
+        async def fetch_embedding(client: httpx.AsyncClient, index: int, text: str) -> None:
+            """Fetch a single embedding."""
+            async with limiter:
+                try:
+                    # Truncate oversized inputs
+                    truncated = text[: self.MAX_EMBED_CHARS] if len(text) > self.MAX_EMBED_CHARS else text
+
+                    response = await client.post(
+                        f"{self.embedding_url}/api/embeddings",
+                        json={"model": model, "prompt": truncated},
+                        timeout=self.timeout,
+                    )
+                    response.raise_for_status()
+                    data = response.json()
+
+                    embedding_list = data.get("embedding", [])
+                    embedding = tuple(embedding_list) if embedding_list else ()
+
+                    if embedding:
+                        self._current_dimensions = len(embedding)
+
+                    results[index] = EmbeddingResponse(
+                        embedding=embedding,
+                        model=model,
+                        provider=self.provider_name,
+                        dimensions=len(embedding),
+                    )
+                except Exception as e:
+                    errors.append(e)
+
+        async with httpx.AsyncClient() as client, trio.open_nursery() as nursery:
+            for i, text in enumerate(texts):
+                nursery.start_soon(fetch_embedding, client, i, text)
+
+        if errors:
+            raise ConnectionError(f"Ollama async batch embed failed: {errors[0]}") from errors[0]
+
+        # Return results in original order
+        return [results[i] for i in range(len(texts))]
+
     def chat(
         self,
         messages: list[dict[str, str]],
@@ -262,9 +420,8 @@ class OllamaProvider(BaseLLMProvider, BaseEmbeddingProvider):
         }
 
         try:
-            response = requests.post(
+            response = self.session.post(
                 f"{self.base_url}/api/chat",
-                headers=self._get_headers(),
                 json=payload,
                 timeout=self.timeout,
             )
@@ -282,3 +439,23 @@ class OllamaProvider(BaseLLMProvider, BaseEmbeddingProvider):
             )
         except requests.RequestException as e:
             raise ConnectionError(f"Ollama chat failed: {e}") from e
+
+    @staticmethod
+    def clear_embedding_cache() -> None:
+        """Clear the embedding cache."""
+        _cached_embedding.cache_clear()
+
+    @staticmethod
+    def embedding_cache_info() -> dict[str, int]:
+        """Get embedding cache statistics."""
+        info = _cached_embedding.cache_info()
+        return {
+            "hits": info.hits,
+            "misses": info.misses,
+            "maxsize": info.maxsize or 0,
+            "currsize": info.currsize,
+        }
+
+
+# Export the EMBEDDING_DIMENSIONS for external use
+EMBEDDING_DIMENSIONS = OllamaProvider.EMBEDDING_DIMENSIONS

ragit/version.py

@@ -2,4 +2,4 @@
 # Copyright RODMENA LIMITED 2025
 # SPDX-License-Identifier: Apache-2.0
 #
-__version__ = "0.7.2"
+__version__ = "0.7.4"

{ragit-0.7.2.dist-info → ragit-0.7.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ragit
-Version: 0.7.2
+Version: 0.7.4
 Summary: Automatic RAG Pattern Optimization Engine
 Author: RODMENA LIMITED
 Maintainer-email: RODMENA LIMITED <info@rodmena.co.uk>
@@ -26,6 +26,8 @@ Requires-Dist: pydantic>=2.0.0
 Requires-Dist: python-dotenv>=1.0.0
 Requires-Dist: scikit-learn>=1.5.0
 Requires-Dist: tqdm>=4.66.0
+Requires-Dist: trio>=0.24.0
+Requires-Dist: httpx>=0.27.0
 Provides-Extra: dev
 Requires-Dist: ragit[test]; extra == "dev"
 Requires-Dist: pytest; extra == "dev"
@@ -443,6 +445,77 @@ print(f"Score: {best.score:.3f}")
 
 The experiment tests different combinations of chunk sizes, overlaps, and retrieval parameters to find what works best for your content.
 
+## Performance Features
+
+Ragit includes several optimizations for production workloads:
+
+### Connection Pooling
+
+`OllamaProvider` uses HTTP connection pooling via `requests.Session()` for faster sequential requests:
+
+```python
+from ragit.providers import OllamaProvider
+
+provider = OllamaProvider()
+
+# All requests reuse the same connection pool
+for text in texts:
+    provider.embed(text, model="mxbai-embed-large")
+
+# Explicitly close when done (optional, auto-closes on garbage collection)
+provider.close()
+```
+
+### Async Parallel Embedding
+
+For large batches, use `embed_batch_async()` with trio for 5-10x faster embedding:
+
+```python
+import trio
+from ragit.providers import OllamaProvider
+
+provider = OllamaProvider()
+
+async def embed_documents():
+    texts = ["doc1...", "doc2...", "doc3...", ...]  # hundreds of texts
+    embeddings = await provider.embed_batch_async(
+        texts,
+        model="mxbai-embed-large",
+        max_concurrent=10  # Adjust based on server capacity
+    )
+    return embeddings
+
+# Run with trio
+results = trio.run(embed_documents)
+```
+
+### Embedding Cache
+
+Repeated embedding calls are cached automatically (2048 entries LRU):
+
+```python
+from ragit.providers import OllamaProvider
+
+provider = OllamaProvider(use_cache=True)  # Default
+
+# First call hits the API
+provider.embed("Hello world", model="mxbai-embed-large")
+
+# Second call returns cached result instantly
+provider.embed("Hello world", model="mxbai-embed-large")
+
+# View cache statistics
+print(OllamaProvider.embedding_cache_info())
+# {'hits': 1, 'misses': 1, 'maxsize': 2048, 'currsize': 1}
+
+# Clear cache if needed
+OllamaProvider.clear_embedding_cache()
+```
+
+### Pre-normalized Embeddings
+
+Vector similarity uses pre-normalized embeddings, making cosine similarity a simple dot product (O(1) per comparison).
+
 ## API Reference
 
 ### Document Loading

{ragit-0.7.2.dist-info → ragit-0.7.4.dist-info}/RECORD

@@ -2,17 +2,17 @@ ragit/__init__.py,sha256=PjQogIWMlydZFWVECqhmxw-X9i7lEXdUTe2XlT6qYUQ,2213
 ragit/assistant.py,sha256=lXjZRUr_WsYLP3XLOktabgfPVyKOZPdREzyL7cSRufk,11251
 ragit/config.py,sha256=uKLchJQHjH8MImZ2OJahDjSzyasFqgrFb9Z4aHxJ7og,1495
 ragit/loaders.py,sha256=keusuPzXPBiLDVj4hKfPCcge-rm-cnzNRk50fGXvTJs,5571
-ragit/version.py,sha256=5FNjmLKNB4z5E3tbMqwASQafUUFSWnaxBRb0EsQPVK8,97
+ragit/version.py,sha256=8-YGrxlAluU3va125prug6u_bleRtoTX3c4m7WfDYNM,97
 ragit/core/__init__.py,sha256=j53PFfoSMXwSbK1rRHpMbo8mX2i4R1LJ5kvTxBd7-0w,100
 ragit/core/experiment/__init__.py,sha256=4vAPOOYlY5Dcr2gOolyhBSPGIUxZKwEkgQffxS9BodA,452
 ragit/core/experiment/experiment.py,sha256=Qh1NJkY9LbKaidRfiI8GOwBZqopjK-MSVBuD_JEgO-k,16582
 ragit/core/experiment/results.py,sha256=KHpN3YSLJ83_JUfIMccRPS-q7LEt0S9p8ehDRawk_4k,3487
 ragit/providers/__init__.py,sha256=iliJt74Lt3mFUlKGfSFW-D0cMonUygY6sRZ6lLjeU7M,435
 ragit/providers/base.py,sha256=MJ8mVeXuGWhkX2XGTbkWIY3cVoTOPr4h5XBXw8rAX2Q,3434
-ragit/providers/ollama.py,sha256=epmXiJ06jCWFERWFMdCNr8FP6WubxdHXj8dUn71WJuQ,9502
+ragit/providers/ollama.py,sha256=fFUziRmpu_L6rnlxi53PnzQ5b34-fEbaU4u1irK1nG0,16055
 ragit/utils/__init__.py,sha256=-UsE5oJSnmEnBDswl-ph0A09Iu8yKNbPhd1-_7Lcb8Y,3051
-ragit-0.7.2.dist-info/licenses/LICENSE,sha256=tAkwu8-AdEyGxGoSvJ2gVmQdcicWw3j1ZZueVV74M-E,11357
-ragit-0.7.2.dist-info/METADATA,sha256=TZ_VuZbe4GKkb8Jo8tC_NgT3hcu7e-uy7S9w6kpURlE,13662
-ragit-0.7.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ragit-0.7.2.dist-info/top_level.txt,sha256=pkPbG7yrw61wt9_y_xcLE2vq2a55fzockASD0yq0g4s,6
-ragit-0.7.2.dist-info/RECORD,,
+ragit-0.7.4.dist-info/licenses/LICENSE,sha256=tAkwu8-AdEyGxGoSvJ2gVmQdcicWw3j1ZZueVV74M-E,11357
+ragit-0.7.4.dist-info/METADATA,sha256=5nEjliSCc-F7X6IgG7OdgkadNjANPw4env04EEBI5J8,15528
+ragit-0.7.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ragit-0.7.4.dist-info/top_level.txt,sha256=pkPbG7yrw61wt9_y_xcLE2vq2a55fzockASD0yq0g4s,6
+ragit-0.7.4.dist-info/RECORD,,