Feature/add sciphi search rag (langroid#350)

* Add sciphi provider * undo readme change * add search rag tool * cleanup * fix flag * add agent-search dependency, pydantic->1.10.13, lints * chat-search small mods * sciphi test -> tests/extras; make agent-search optional --------- Co-authored-by: Owen <ocolegro@gmail.com>
dragon28 · Jan 8, 2024 · 151a11f · 151a11f
1 parent b00ba1d
commit 151a11f
Show file tree

Hide file tree

Showing 3 changed files with 126 additions and 15 deletions.
diff --git a/examples/basic/chat-search.py b/examples/basic/chat-search.py
@@ -1,28 +1,46 @@
 """
-This is a basic example of a chatbot that uses the GoogleSearchTool:
-when the LLM doesn't know the answer to a question, it will use the tool to
+This is a basic example of a chatbot that uses the GoogleSearchTool
+or SciPhiSearchRAGTool to answer questions.
+When the LLM doesn't know the answer to a question, it will use the tool to
 search the web for relevant results, and then use the results to answer the
 question.
 
-NOTE: running this example requires setting the GOOGLE_API_KEY and GOOGLE_CSE_ID
+Run like this:
+
+python3 examples/basic/chat-search.py
+
+You can specify which search provider to use with this optional flag:
+
+-p or --provider: google or sciphi (default: google)
+
+
+NOTE:
+(a) If using Google Search, you must have GOOGLE_API_KEY and GOOGLE_CSE_ID
 environment variables in your `.env` file, as explained in the
 [README](https://github.com/langroid/langroid#gear-installation-and-setup).
+
+(b) Alternatively, you can use the SciPhiSearchRAGTool, you need to have the
+SCIPHI_API_KEY environment variable in your `.env` file.
+See here for more info: https://www.sciphi.ai/
+This tool requires installing langroid with the `sciphi` extra, e.g.
+`pip install langroid[sciphi]` or `poetry add langroid[sciphi]`
+(it installs the `agent-search` package from pypi).
 """
 
 import typer
+from dotenv import load_dotenv
+from pydantic import BaseSettings
 from rich import print
 from rich.prompt import Prompt
-from pydantic import BaseSettings
-from dotenv import load_dotenv
 
 from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
 from langroid.agent.task import Task
 from langroid.agent.tools.google_search_tool import GoogleSearchTool
+from langroid.agent.tools.sciphi_search_rag_tool import SciPhiSearchRAGTool
 from langroid.language_models.openai_gpt import OpenAIGPTConfig
-from langroid.utils.configuration import set_global, Settings
+from langroid.utils.configuration import Settings, set_global
 from langroid.utils.logging import setup_colored_logging
 
-
 app = typer.Typer()
 
 setup_colored_logging()
@@ -50,6 +68,7 @@
 
 class CLIOptions(BaseSettings):
     model: str = ""
+    provider: str = "google"
 
     class Config:
         extra = "forbid"
@@ -59,9 +78,9 @@ class Config:
 def chat(opts: CLIOptions) -> None:
     print(
         """
-        [blue]Welcome to the Google Search chatbot!
+        [blue]Welcome to the Web Search chatbot!
         I will try to answer your questions, relying on (summaries of links from) 
-        Google Search when needed.
+        Search when needed.
         
         Enter x or q to quit at any point.
         """
@@ -88,14 +107,25 @@ def chat(opts: CLIOptions) -> None:
         vecdb=None,
     )
     agent = ChatAgent(config)
-    agent.enable_message(GoogleSearchTool)
+
+    match opts.provider:
+        case "google":
+            search_tool_class = GoogleSearchTool
+        case "sciphi":
+            search_tool_class = SciPhiSearchRAGTool
+        case _:
+            raise ValueError(f"Unsupported provider {opts.provider} specified.")
+
+    agent.enable_message(search_tool_class)
+    search_tool_handler_method = search_tool_class.default_value("request")
+
     task = Task(
         agent,
-        system_message="""
+        system_message=f"""
         You are a helpful assistant. You will try your best to answer my questions.
         If you cannot answer from your own knowledge, you can use up to 5 
-        results from the `web_search` tool/function-call to help you with 
-        answering the question.
+        results from the {search_tool_handler_method} tool/function-call to help 
+        you with answering the question.
         Be very concise in your responses, use no more than 1-2 sentences.
         When you answer based on a web search, First show me your answer, 
         and then show me the SOURCE(s) and EXTRACT(s) to justify your answer,
@@ -120,6 +150,9 @@ def chat(opts: CLIOptions) -> None:
 def main(
     debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
     model: str = typer.Option("", "--model", "-m", help="model name"),
+    provider: str = typer.Option(
+        "google", "--provider", "-p", help="search provider name (Google, SciPhi)"
+    ),
     no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
     nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
     cache_type: str = typer.Option(
@@ -134,7 +167,7 @@ def main(
             cache_type=cache_type,
         )
     )
-    opts = CLIOptions(model=model)
+    opts = CLIOptions(model=model, provider=provider)
     chat(opts)
 
 

diff --git a/langroid/agent/tools/sciphi_search_rag_tool.py b/langroid/agent/tools/sciphi_search_rag_tool.py
@@ -0,0 +1,76 @@
+"""
+A tool which returns a Search RAG response from the SciPhi API.
+their titles, links, summaries. Since the tool is stateless (i.e. does not need
+access to agent state), it can be enabled for any agent, without having to define a
+special method inside the agent: `agent.enable_message(SciPhiSearchRAGTool)`
+
+Example return output appears as follows below:
+
+<-- Query -->
+```
+Find 3 results on the internet about the LK-99 superconducting material.
+``
+
+<-- Response (compressed for this example)-->
+```
+[ result1 ]
+
+[ result2 ]
+
+[ result3 ]
+
+```
+
+NOTE: Using this tool requires getting an API key from sciphi.ai.
+Setup is as simple as shown below
+# Get a free API key at https://www.sciphi.ai/account
+# export SCIPHI_API_KEY=$MY_SCIPHI_API_KEY before running the agent
+# OR add SCIPHI_API_KEY=$MY_SCIPHI_API_KEY to your .env file
+
+This tool requires installing langroid with the `sciphi` extra, e.g.
+`pip install langroid[sciphi]` or `poetry add langroid[sciphi]`
+(it installs the `agent-search` package from pypi).
+
+For more information, please refer to the official docs:
+https://agent-search.readthedocs.io/en/latest/
+"""
+
+try:
+    from agent_search import SciPhi
+except ImportError:
+    raise ImportError(
+        "You are attempting to use the `agent-search` library;"
+        "To use it, please install langroid with the `sciphi` extra, e.g. "
+        "`pip install langroid[sciphi]` or `poetry add langroid[sciphi]` "
+        "(it installs the `agent-search` package from pypi)."
+    )
+
+from langroid.agent.tool_message import ToolMessage
+
+
+class SciPhiSearchRAGTool(ToolMessage):
+    request: str = "web_search_rag"
+    purpose: str = """
+            To search the web with provider <search_provider> and 
+            return a response summary with llm model <llm_model> the given <query>. 
+            """
+    query: str
+    search_provider: str = "bing"  # bing or agent-search
+    include_related_queries: bool = True
+    llm_model: str = "SciPhi/Sensei-7B-V1"
+    recursive_mode: bool = True
+
+    def handle(self) -> str:
+        rag_response = SciPhi().get_search_rag_response(
+            query=self.query,
+            search_provider=self.search_provider,
+            llm_model=self.llm_model,
+        )
+        result = rag_response["response"]
+        if self.include_related_queries:
+            result = (
+                f"### RAG Response:\n{result}\n\n"
+                + "### Related Queries:\n"
+                + "\n".join(rag_response["related_queries"])
+            )
+        return result  # type: ignore
diff --git a/pyproject.toml b/pyproject.toml
@@ -47,7 +47,7 @@ types-requests = "^2.31.0.1"
 pyparsing = "^3.0.9"
 nltk = "^3.8.1"
 qdrant-client = "^1.7.0"
-pydantic = "1.10.11"
+pydantic = "1.10.13"
 pypdf = "^3.12.2"
 momento = "^1.10.2"
 pandas = "^2.0.3"
@@ -81,6 +81,7 @@ scrapy = "^2.11.0"
 async-generator = "^1.10"
 lancedb = "^0.4.1"
 pytest-redis = "^3.0.2"
+agent-search = {version = "^0.0.7", optional = true}
 python-docx = "^1.1.0"
 
 [tool.poetry.extras]
@@ -89,6 +90,7 @@ hf-embeddings = ["sentence-transformers", "torch"]
 postgres = ["psycopg2", "pytest-postgresql"]
 mysql = ["pymysql", "pytest-mysql"]
 litellm = ["litellm"]
+sciphi = ["agent-search"]
 
 
 [tool.poetry.group.dev.dependencies]