Initial commit
This commit is contained in:
yichuan520030910320
179
packages/leann-core/src/leann/api.py
Normal file
179
packages/leann-core/src/leann/api.py
Normal file
@@ -0,0 +1,179 @@
|
||||
from .registry import BACKEND_REGISTRY
|
||||
from .interface import LeannBackendFactoryInterface
|
||||
from typing import List, Dict, Any, Optional
|
||||
import numpy as np
|
||||
import os
|
||||
import json
|
||||
from pathlib import Path
|
||||
import openai # Import openai library
|
||||
|
||||
# 一个辅助函数，用于临时计算 embedding
|
||||
def _compute_embeddings(chunks: List[str], model_name: str) -> np.ndarray:
|
||||
try:
|
||||
from sentence_transformers import SentenceTransformer
|
||||
model = SentenceTransformer(model_name)
|
||||
print(f"INFO: Computing embeddings for {len(chunks)} chunks using '{model_name}'...")
|
||||
embeddings = model.encode(chunks, show_progress_bar=True)
|
||||
return np.asarray(embeddings, dtype=np.float32)
|
||||
except ImportError:
|
||||
print("WARNING: sentence-transformers not installed. Falling back to random embeddings.")
|
||||
# 如果没有安装，则生成随机向量用于测试
|
||||
# TODO: 应该从一个固定的地方获取维度信息
|
||||
return np.random.rand(len(chunks), 768).astype(np.float32)
|
||||
|
||||
|
||||
class LeannBuilder:
|
||||
"""
|
||||
负责构建 Leann 索引的上层 API。
|
||||
它协调 embedding 计算和后端索引构建。
|
||||
"""
|
||||
def __init__(self, backend_name: str, embedding_model: str = "sentence-transformers/all-mpnet-base-v2", **backend_kwargs):
|
||||
self.backend_name = backend_name
|
||||
self.backend_factory = BACKEND_REGISTRY.get(backend_name)
|
||||
if self.backend_factory is None:
|
||||
raise ValueError(f"Backend '{backend_name}' not found or not registered.")
|
||||
|
||||
self.embedding_model = embedding_model
|
||||
self.backend_kwargs = backend_kwargs
|
||||
self.chunks: List[Dict[str, Any]] = []
|
||||
print(f"INFO: LeannBuilder initialized with '{backend_name}' backend.")
|
||||
|
||||
def add_text(self, text: str, metadata: Optional[Dict[str, Any]] = None):
|
||||
# 简单的分块逻辑
|
||||
self.chunks.append({"text": text, "metadata": metadata or {}})
|
||||
|
||||
def build_index(self, index_path: str):
|
||||
if not self.chunks:
|
||||
raise ValueError("No chunks added. Use add_text() first.")
|
||||
|
||||
# 1. 计算 embedding (这是 leann-core 的职责)
|
||||
texts_to_embed = [c["text"] for c in self.chunks]
|
||||
embeddings = _compute_embeddings(texts_to_embed, self.embedding_model)
|
||||
|
||||
# 2. 创建 builder 实例并构建索引
|
||||
builder_instance = self.backend_factory.builder(**self.backend_kwargs)
|
||||
builder_instance.build(embeddings, index_path, **self.backend_kwargs)
|
||||
|
||||
# 3. 保存 leann 特有的元数据（不包含向量）
|
||||
index_dir = Path(index_path).parent
|
||||
leann_meta_path = index_dir / f"{Path(index_path).name}.meta.json"
|
||||
|
||||
meta_data = {
|
||||
"version": "0.1.0",
|
||||
"backend_name": self.backend_name,
|
||||
"embedding_model": self.embedding_model,
|
||||
"num_chunks": len(self.chunks),
|
||||
"chunks": self.chunks,
|
||||
}
|
||||
with open(leann_meta_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(meta_data, f, indent=2)
|
||||
print(f"INFO: Leann metadata saved to {leann_meta_path}")
|
||||
|
||||
|
||||
class LeannSearcher:
|
||||
"""
|
||||
负责加载索引并执行检索的上层 API。
|
||||
"""
|
||||
def __init__(self, index_path: str, **backend_kwargs):
|
||||
leann_meta_path = Path(index_path).parent / f"{Path(index_path).name}.meta.json"
|
||||
if not leann_meta_path.exists():
|
||||
raise FileNotFoundError(f"Leann metadata file not found at {leann_meta_path}. Was the index built with LeannBuilder?")
|
||||
|
||||
with open(leann_meta_path, 'r', encoding='utf-8') as f:
|
||||
self.meta_data = json.load(f)
|
||||
|
||||
backend_name = self.meta_data['backend_name']
|
||||
self.embedding_model = self.meta_data['embedding_model']
|
||||
|
||||
backend_factory = BACKEND_REGISTRY.get(backend_name)
|
||||
if backend_factory is None:
|
||||
raise ValueError(f"Backend '{backend_name}' (from index file) not found or not registered.")
|
||||
|
||||
# 创建 searcher 实例
|
||||
self.backend_impl = backend_factory.searcher(index_path, **backend_kwargs)
|
||||
print(f"INFO: LeannSearcher initialized with '{backend_name}' backend using index '{index_path}'.")
|
||||
|
||||
def search(self, query: str, top_k: int = 5, **search_kwargs):
|
||||
query_embedding = _compute_embeddings([query], self.embedding_model)
|
||||
|
||||
# 委托给后端的 search 方法
|
||||
results = self.backend_impl.search(query_embedding, top_k, **search_kwargs)
|
||||
|
||||
# 丰富返回结果，加入原始文本和元数据
|
||||
enriched_results = []
|
||||
for label, dist in zip(results['labels'][0], results['distances'][0]):
|
||||
if label < len(self.meta_data['chunks']):
|
||||
chunk_info = self.meta_data['chunks'][label]
|
||||
enriched_results.append({
|
||||
"id": label,
|
||||
"score": dist,
|
||||
"text": chunk_info['text'],
|
||||
"metadata": chunk_info['metadata']
|
||||
})
|
||||
return enriched_results
|
||||
|
||||
|
||||
class LeannChat:
|
||||
"""
|
||||
封装了 Searcher 和 LLM 的对话式 RAG 接口。
|
||||
"""
|
||||
def __init__(self, index_path: str, backend_name: Optional[str] = None, llm_model: str = "gpt-4o", **kwargs):
|
||||
# 如果用户没有指定后端，尝试从索引元数据中读取
|
||||
if backend_name is None:
|
||||
leann_meta_path = Path(index_path).parent / f"{Path(index_path).name}.meta.json"
|
||||
if not leann_meta_path.exists():
|
||||
raise FileNotFoundError(f"Leann metadata file not found at {leann_meta_path}.")
|
||||
with open(leann_meta_path, 'r', encoding='utf-8') as f:
|
||||
meta_data = json.load(f)
|
||||
backend_name = meta_data['backend_name']
|
||||
|
||||
self.searcher = LeannSearcher(index_path, **kwargs)
|
||||
self.llm_model = llm_model
|
||||
self.openai_client = None # Lazy load
|
||||
|
||||
def _get_openai_client(self):
|
||||
if self.openai_client is None:
|
||||
api_key = os.getenv("OPENAI_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError("OPENAI_API_KEY environment variable not set.")
|
||||
self.openai_client = openai.OpenAI(api_key=api_key)
|
||||
return self.openai_client
|
||||
|
||||
def ask(self, question: str, **kwargs):
|
||||
# 1. 检索
|
||||
results = self.searcher.search(question, top_k=3, **kwargs)
|
||||
context = "\n\n".join([r['text'] for r in results])
|
||||
|
||||
# 2. 构建 Prompt
|
||||
prompt = f"Context:\n{context}\n\nQuestion: {question}\n\nAnswer:"
|
||||
|
||||
# 3. 调用 LLM
|
||||
print(f"DEBUG: Calling LLM with prompt: {prompt[:200]}...")
|
||||
try:
|
||||
client = self._get_openai_client()
|
||||
response = client.chat.completions.create(
|
||||
model=self.llm_model,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant that answers questions based on the provided context."},
|
||||
{"role": "user", "content": prompt}
|
||||
]
|
||||
)
|
||||
return response.choices[0].message.content
|
||||
except Exception as e:
|
||||
print(f"ERROR: Failed to call OpenAI API: {e}")
|
||||
return f"Error: Could not get a response from the LLM. {e}"
|
||||
|
||||
def start_interactive(self):
|
||||
print("\nLeann Chat started (type 'quit' to exit)")
|
||||
while True:
|
||||
try:
|
||||
user_input = input("You: ").strip()
|
||||
if user_input.lower() in ['quit', 'exit']:
|
||||
break
|
||||
if not user_input:
|
||||
continue
|
||||
response = self.ask(user_input)
|
||||
print(f"Leann: {response}")
|
||||
except (KeyboardInterrupt, EOFError):
|
||||
print("\nGoodbye!")
|
||||
break
|
||||
59
packages/leann-core/src/leann/interface.py
Normal file
59
packages/leann-core/src/leann/interface.py
Normal file
@@ -0,0 +1,59 @@
|
||||
from abc import ABC, abstractmethod
|
||||
import numpy as np
|
||||
from typing import Dict, Any
|
||||
|
||||
class LeannBackendBuilderInterface(ABC):
|
||||
"""用于构建索引的后端接口"""
|
||||
|
||||
@abstractmethod
|
||||
def build(self, data: np.ndarray, index_path: str, **kwargs) -> None:
|
||||
"""构建索引
|
||||
|
||||
Args:
|
||||
data: 向量数据 (N, D)
|
||||
index_path: 索引保存路径
|
||||
**kwargs: 后端特定的构建参数
|
||||
"""
|
||||
pass
|
||||
|
||||
class LeannBackendSearcherInterface(ABC):
|
||||
"""用于搜索的后端接口"""
|
||||
|
||||
@abstractmethod
|
||||
def __init__(self, index_path: str, **kwargs):
|
||||
"""初始化搜索器
|
||||
|
||||
Args:
|
||||
index_path: 索引文件路径
|
||||
**kwargs: 后端特定的加载参数
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def search(self, query: np.ndarray, top_k: int, **kwargs) -> Dict[str, Any]:
|
||||
"""搜索最近邻
|
||||
|
||||
Args:
|
||||
query: 查询向量 (1, D) 或 (B, D)
|
||||
top_k: 返回的最近邻数量
|
||||
**kwargs: 搜索参数
|
||||
|
||||
Returns:
|
||||
{"labels": [...], "distances": [...]}
|
||||
"""
|
||||
pass
|
||||
|
||||
class LeannBackendFactoryInterface(ABC):
|
||||
"""后端工厂接口"""
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def builder(**kwargs) -> LeannBackendBuilderInterface:
|
||||
"""创建 Builder 实例"""
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def searcher(index_path: str, **kwargs) -> LeannBackendSearcherInterface:
|
||||
"""创建 Searcher 实例"""
|
||||
pass
|
||||
12
packages/leann-core/src/leann/registry.py
Normal file
12
packages/leann-core/src/leann/registry.py
Normal file
@@ -0,0 +1,12 @@
|
||||
# packages/leann-core/src/leann/registry.py
|
||||
|
||||
# 全局的后端注册表字典
|
||||
BACKEND_REGISTRY = {}
|
||||
|
||||
def register_backend(name: str):
|
||||
"""一个用于注册新后端类的装饰器。"""
|
||||
def decorator(cls):
|
||||
print(f"INFO: Registering backend '{name}'")
|
||||
BACKEND_REGISTRY[name] = cls
|
||||
return cls
|
||||
return decorator
|
||||
Reference in New Issue
Block a user