custom retriever

.gitignore

@@ -38,7 +38,9 @@ venv*
 phienv*
 aienv*
 
-# ignore scratch dir
+# ignore tmp dirs
 scratch
+junk
+tmp
 
 .ipynb_checkpoints

cookbook/knowledge/README.md

@@ -25,12 +25,14 @@ pip install -U openai phidata
 phi start cookbook/knowledge/resources.py -y
 ```
 
-4. PDF URL Knowledge Base
+4. Test knowledge cookbooks
+
+Eg: PDF URL Knowledge Base
 
 - Install libraries
 
 ```shell
-pip install -U pypdf sqlalchemy pgvector psycopg-binary bs4
+pip install -U pypdf sqlalchemy pgvector 'psycopg[binary]' bs4
 ```
 
 - Run the PDF URL script

cookbook/knowledge/manual_references.py

@@ -0,0 +1,36 @@
+import json
+from typing import List, Optional
+
+from phi.assistant import Assistant
+from phi.document import Document
+from phi.knowledge.pdf import PDFUrlKnowledgeBase
+from phi.vectordb.pgvector import PgVector2
+
+from resources import vector_db  # type: ignore
+
+knowledge_base = PDFUrlKnowledgeBase(
+    urls=["https://phi-public.s3.amazonaws.com/recipes/ThaiRecipes.pdf"],
+    vector_db=PgVector2(collection="recipes", db_url=vector_db.get_db_connection_local()),
+)
+# Comment out after first run
+# knowledge_base.load(recreate=False)
+
+
+def custom_references_function(query: str, **kwargs) -> Optional[str]:
+    """Return a list of references from the knowledge base"""
+    print(f"-*- Searching for references for query: {query}")
+    relevant_docs: List[Document] = knowledge_base.search(query=query, num_documents=5)
+    if len(relevant_docs) == 0:
+        return None
+
+    return json.dumps([doc.to_dict() for doc in relevant_docs], indent=2)
+
+
+assistant = Assistant(
+    knowledge_base=knowledge_base,
+    # Generate references using a custom function.
+    references_function=custom_references_function,
+    # Adds references to the prompt.
+    add_references_to_prompt=True,
+)
+assistant.print_response("How to make Thai curry?", markdown=True)

phi/assistant/__init__.py

@@ -1 +1,12 @@
-from phi.assistant.assistant import Assistant, AssistantRun, AssistantMemory, AssistantStorage, AssistantKnowledge
+from phi.assistant.assistant import (
+    Assistant,
+    AssistantRun,
+    AssistantMemory,
+    AssistantStorage,
+    AssistantKnowledge,
+    LLMTask,
+    Task,
+    Function,
+    Tool,
+    Toolkit,
+)

phi/task/llm/llm_task.py

@@ -388,7 +388,7 @@ def get_references_from_knowledge_base(self, query: str, num_documents: Optional
         """Return a list of references from the knowledge base"""
 
         if self.references_function is not None:
-            reference_kwargs = {"task": self, "query": query}
+            reference_kwargs = {"task": self, "query": query, "num_documents": num_documents}
             return remove_indent(self.references_function(**reference_kwargs))
 
         if self.knowledge_base is None:

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "phidata"
-version = "2.3.53"
+version = "2.3.54"
 description = "Build AI Assistants using function calling"
 requires-python = ">=3.8"
 readme = "README.md"
@@ -71,7 +71,7 @@ check_untyped_defs = true
 no_implicit_optional = true
 warn_unused_configs = true
 plugins = ["pydantic.mypy"]
-exclude = ["phienv*", "aienv*", "scratch*"]
+exclude = ["phienv*", "aienv*", "scratch*", "junk*", "tmp*"]
 
 [[tool.mypy.overrides]]
 module = [