Update retriever.py

retriever.py

@@ -1,53 +1,65 @@
-from smolagents import Tool
-from langchain_community.retrievers import BM25Retriever
-from langchain.docstore.document import Document
-import datasets
-
-
-class GuestInfoRetrieverTool(Tool):
-    name = "guest_info_retriever"
-    description = "Retrieves detailed information about gala guests based on their name or relation."
-    inputs = {
-        "query": {
-            "type": "string",
-            "description": "The name or relation of the guest you want information about."
-        }
-    }
-    output_type = "string"
-
-    def __init__(self, docs):
-        self.is_initialized = False
-        self.retriever = BM25Retriever.from_documents(docs)
+# from smolagents import Tool
+# from langchain_community.retrievers import BM25Retriever
+# from langchain.docstore.document import Document
+# import datasets
+
+
+# class GuestInfoRetrieverTool(Tool):
+#     name = "guest_info_retriever"
+#     description = "Retrieves detailed information about gala guests based on their name or relation."
+#     inputs = {
+#         "query": {
+#             "type": "string",
+#             "description": "The name or relation of the guest you want information about."
+#         }
+#     }
+#     output_type = "string"
+
+#     def __init__(self, docs):
+#         self.is_initialized = False
+#         self.retriever = BM25Retriever.from_documents(docs)
        
 
-    def forward(self, query: str):
-        results = self.retriever.get_relevant_documents(query)
-        if results:
-            return "\n\n".join([doc.page_content for doc in results[:3]])
-        else:
-            return "No matching guest information found."
+#     def forward(self, query: str):
+#         results = self.retriever.get_relevant_documents(query)
+#         if results:
+#             return "\n\n".join([doc.page_content for doc in results[:3]])
+#         else:
+#             return "No matching guest information found."
+
 
+# def load_guest_dataset():
+#     # Load the dataset
+#     guest_dataset = datasets.load_dataset("agents-course/unit3-invitees", split="train")
 
-def load_guest_dataset():
-    # Load the dataset
-    guest_dataset = datasets.load_dataset("agents-course/unit3-invitees", split="train")
+#     # Convert dataset entries into Document objects
+#     docs = [
+#         Document(
+#             page_content="\n".join([
+#                 f"Name: {guest['name']}",
+#                 f"Relation: {guest['relation']}",
+#                 f"Description: {guest['description']}",
+#                 f"Email: {guest['email']}"
+#             ]),
+#             metadata={"name": guest["name"]}
+#         )
+#         for guest in guest_dataset
+#     ]
 
-    # Convert dataset entries into Document objects
-    docs = [
-        Document(
-            page_content="\n".join([
-                f"Name: {guest['name']}",
-                f"Relation: {guest['relation']}",
-                f"Description: {guest['description']}",
-                f"Email: {guest['email']}"
-            ]),
-            metadata={"name": guest["name"]}
-        )
-        for guest in guest_dataset
-    ]
+#     # Return the tool
+#     return GuestInfoRetrieverTool(docs)
 
-    # Return the tool
-    return GuestInfoRetrieverTool(docs)
 
+from langchain_community.vectorstores import Chroma
+from langchain_openai import OpenAIEmbeddings
 
+def get_retriever():
+    """创建并返回检索器"""
+    embeddings = OpenAIEmbeddings()
+    vectorstore = Chroma(
+        embedding_function=embeddings,
+        persist_directory="./chroma_db",
+        collection_name="rag_docs"
+    )
+    return vectorstore.as_retriever(search_kwargs={"k": 5})