Merge pull request #3 from hemslo/generic-search

Generic search

app/agents/search_agent.py

@@ -0,0 +1,37 @@
+from langchain.agents import AgentExecutor
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain.globals import set_debug
+from .util import create_ollama_functions_agent
+from ..dependencies.ollama_functions_model import ollama_functions_model
+from ..tools.duckduckgo_search import duckduckgo_search
+from typing import Any
+from pydantic import BaseModel
+
+prompt = ChatPromptTemplate.from_messages(
+    [
+        (
+            "user",
+            "{input}",
+        ),
+        MessagesPlaceholder(variable_name="agent_scratchpad"),
+    ]
+)
+
+class Input(BaseModel):
+    input: str
+
+
+class Output(BaseModel):
+    output: Any
+
+tools = [duckduckgo_search]
+llm = ollama_functions_model
+agent = create_ollama_functions_agent(llm, tools, prompt)
+
+search_engine_agent_executor = AgentExecutor(
+    agent=agent,
+    tools=tools,
+).with_types(
+    input_type=Input,
+    output_type=Output,
+)

app/graph.py

@@ -15,6 +15,7 @@
 from app.dependencies.openai_chat_model import openai_chat_model
 from app.tools.random_number import random_number
 from app.tools.random_select import random_select
+from app.tools.duckduckgo_search import duckduckgo_search
 
 
 # https://github.com/langchain-ai/langgraph/blob/main/examples/multi_agent/agent_supervisor.ipynb
@@ -55,6 +56,10 @@ class AgentState(TypedDict):
     "RandomSelect": {
         "tools": [random_select],
         "system_prompt": "You are a random selector.",
+    },
+    "GeneriaSearch": {
+        "tools": [duckduckgo_search],
+        "system_prompt": "You are a search engine for generic questions.",
     }
 }
 

app/server.py

@@ -3,6 +3,7 @@
 from langserve import add_routes
 
 from app.agents.random_number import random_number_agent_executor
+from app.agents.search_agent import search_engine_agent_executor
 from app.chains.extraction import extraction_chain
 from app.chains.supervisor import build_supervisor_chain
 from app.dependencies.ollama_chat_model import ollama_chat_model
@@ -35,6 +36,12 @@ async def redirect_root_to_docs():
     path="/random_number",
 )
 
+add_routes(
+    app,
+    search_engine_agent_executor,
+    path="/generic",
+)
+
 add_routes(
     app,
     build_supervisor_chain(["SlackSummarizer", "SlackSearcher", "WebRCA"]),

app/tools/duckduckgo_search.py

@@ -0,0 +1,15 @@
+from langchain_community.tools import DuckDuckGoSearchResults
+from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
+
+
+from langchain_core.tools import tool
+
+
+@tool()
+def duckduckgo_search():
+    """
+    Get search result from Duckduckgo
+    """
+    wrapper = DuckDuckGoSearchAPIWrapper(max_results=10)
+    search = DuckDuckGoSearchResults(api_wrapper=wrapper)
+    return search