73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
from dataclasses import dataclass, field
|
|
from typing import Type
|
|
import tyro
|
|
from mcp.server.fastmcp import FastMCP
|
|
from typing import List
|
|
import tyro
|
|
|
|
from langchain_community.vectorstores import FAISS
|
|
from langchain_core.documents.base import Document
|
|
|
|
from lang_agent.rag.emb import QwenEmbeddings
|
|
from lang_agent.config import InstantiateConfig
|
|
|
|
|
|
mcp = FastMCP("Rag")
|
|
|
|
@tyro.conf.configure(tyro.conf.SuppressFixed)
|
|
@dataclass
|
|
class SimpleRagConfig(InstantiateConfig):
|
|
_target: Type = field(default_factory=lambda: SimpleRag)
|
|
|
|
model_name:str = "text-embedding-v4"
|
|
"""embedding model name"""
|
|
|
|
api_key:str = "wrong-key"
|
|
"""api_key for model; for generic text splitting; give a wrong key"""
|
|
|
|
folder_path:str = "assets/xiaozhan_emb"
|
|
"""path to local database"""
|
|
|
|
|
|
class SimpleRag:
|
|
def __init__(self, config:SimpleRagConfig):
|
|
self.config = config
|
|
self.emb = QwenEmbeddings(self.config.api_key,
|
|
self.config.model_name)
|
|
self.vec_store = FAISS.load_local(
|
|
folder_path=self.config.folder_path,
|
|
embeddings=self.emb,
|
|
allow_dangerous_deserialization=True # Required for LangChain >= 0.1.1
|
|
)
|
|
|
|
# self.retriever = self.vec_store.as_retriever(search_kwargs={"k":3})
|
|
|
|
@mcp.tool()
|
|
def retrieve(self, query:str):
|
|
"""
|
|
检索与给定查询相关的文档，并将其序列化为字符串格式。
|
|
参数:
|
|
query (str): 用户输入的查询字符串。
|
|
返回:
|
|
Tuple[str, List[Document]]:
|
|
- 序列化后的文档内容字符串，每个文档包含来源和内容。
|
|
- 检索到的 Document 对象列表。
|
|
该工具用于基于向量存储检索相关文档，适用于问答和知识检索场景。
|
|
|
|
用例示例:
|
|
1. 用户询问“推荐一些辣味食物”，系统会检索并返回相关的辣味美食推荐文档。
|
|
2. 用户搜索“适合夏天的清爽饮品”，系统会检索并返回相关饮品推荐及其来源信息。
|
|
"""
|
|
retrieved_docs:List[Document] = self.vec_store.search(query, search_kwargs={"k":3})
|
|
serialized = "\n\n".join(
|
|
(f"Source: {doc.metadata}\nContent: {doc.page_content}")
|
|
for doc in retrieved_docs
|
|
)
|
|
return serialized, retrieved_docs
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# config = tyro.cli(SimpleRagConfig)
|
|
config = SimpleRagConfig()
|
|
rag:SimpleRag = config.setup()
|
|
mcp.run(transport="stdio") |