feat: solution for 'Агент с RAG-памятью'

src/embeddings.py

@@ -1,18 +1,75 @@
 """
-Embeddings module that provides an OllamaEmbeddings instance.
+Embeddings module using Ollama.
+
+Provides a simple caching layer and a function to embed text using Ollama's
+embedding endpoint.  No OpenAI services are used.
 """
 
+import json
 import os
-from langchain_community.embeddings import OllamaEmbeddings
+from typing import List, Dict
 
-def get_embedding_model() -> OllamaEmbeddings:
-    """
-    Returns an OllamaEmbeddings instance configured with the model name
-    specified by the OLLAMA_MODEL environment variable or defaults to
-    'mistral'.
+import ollama
+import numpy as np
 
-    Returns:
-        OllamaEmbeddings: The embedding model instance.
+# Cache to avoid repeated calls for the same text
+_EMBED_CACHE: Dict[str, List[float]] = {}
+
+
+def embed(text: str, model: str = "llama2") -> List[float]:
     """
-    model_name = os.getenv("OLLAMA_MODEL", "mistral")
-    return OllamaEmbeddings(model=model_name)
+    Generate an embedding vector for the given text using Ollama.
+
+    Parameters
+    ----------
+    text : str
+        The text to embed.
+    model : str, optional
+        The Ollama model to use for embeddings. Defaults to "llama2".
+
+    Returns
+    -------
+    List[float]
+        The embedding vector.
+    """
+    if text in _EMBED_CACHE:
+        return _EMBED_CACHE[text]
+
+    # Ollama expects a dict with "model" and "prompt"
+    payload = {"model": model, "prompt": text}
+    try:
+        response = ollama.embeddings(payload)
+    except Exception as exc:
+        raise RuntimeError(f"Failed to get embeddings from Ollama: {exc}") from exc
+
+    # Ollama returns a dict with "embedding" key
+    embedding = response.get("embedding")
+    if embedding is None:
+        raise ValueError("Ollama response missing 'embedding' field")
+
+    _EMBED_CACHE[text] = embedding
+    return embedding
+
+
+def cosine_similarity(vec1: List[float], vec2: List[float]) -> float:
+    """
+    Compute cosine similarity between two vectors.
+
+    Parameters
+    ----------
+    vec1, vec2 : List[float]
+        Input vectors.
+
+    Returns
+    -------
+    float
+        Cosine similarity score.
+    """
+    v1 = np.array(vec1)
+    v2 = np.array(vec2)
+    dot = np.dot(v1, v2)
+    norm1 = np.linalg.norm(v1)
+    norm2 = np.linalg.norm(v2)
+    if norm1 == 0 or norm2 == 0:
+        return 0.0
+    return dot / (norm1 * norm2)