222

embedding_ollama.py

@@ -5,15 +5,18 @@ from typing import List
 class OllamaEmbeddings:
     """
     Ollama 本地服务提供的 Embedding 接口，
-    本需求里我们最终拼出形如: http://localhost:11434/api/embed
+    最终拼出形如: http://localhost:11434/api/embed
     即 base_url + "/embed"
     """
 
     def __init__(self, model_name: str, base_url: str):
         self.model_name = model_name
-        self.base_url = base_url  # 这里应形如 http://localhost:11434/api (不再含 /v1)
+        self.base_url = base_url  # 这里一般形如 http://localhost:11434/api (而非 /v1)
 
     def embed(self, texts: List[str]) -> List[List[float]]:
+        """
+        批量将多段文本转换为embedding向量
+        """
         embeddings = []
         for text in texts:
             embeddings.append(self.embed_single_document(text))
@@ -21,13 +24,9 @@ class OllamaEmbeddings:
     
     def embed_documents(self, texts: List[str]) -> List[List[float]]:
         """
-        将多段文本转换为向量列表
+        兼容langchain的接口写法
         """
-        embeddings = []
-        for text in texts:
-            emb = self.embed_single_document(text)
-            embeddings.append(emb)
-        return embeddings
+        return self.embed(texts)
 
     def embed_query(self, query: str) -> List[float]:
         """
@@ -38,7 +37,6 @@ class OllamaEmbeddings:
     def embed_single_document(self, text: str) -> List[float]:
         """
         调用 Ollama 本地服务接口，获取文本的 embedding。
-        这里统一改为请求:  [base_url]/embed
         """
         url = f"{self.base_url}/embed"
         data = {
@@ -49,6 +47,8 @@ class OllamaEmbeddings:
             response = requests.post(url, json=data)
             response.raise_for_status()
             result = response.json()
+            if "embedding" not in result:
+                raise ValueError("No 'embedding' field in Ollama response.")
             return result["embedding"]
         except requests.exceptions.RequestException as e:
             raise Exception(f"Ollama embeddings request error: {e}")