Can I Use the Same Pipeline for Multiple Agents with Different Knowledge Bases?

[learner524]

Can I use the same pipeline for multiple agents with different knowledge bases? For example, suppose I have an agent with a team of two sub-agents: one for Google Search and another for searching from a knowledge base. I have multiple applications that need to use this single pipeline. Is it possible to pass the application name with each request so that the system queries and retrieves responses from the specific collection related to that application? Each application has its own separate collection.

I tried testing a very simple approach where all components are initialized on each call, but it is not working as expected. When I run this locally for a single query, the retrieval of the context and the question-response process works pretty well. However, when I transform this into an API and run it, the same queries produce different (and incorrect) answers. Additionally, using this approach results in high processing time for each query. How can I reduce the processing time as well?

from fastapi import FastAPI, Form
from phi.agent import Agent
from phi.model.openai import OpenAIChat
from phi.storage.agent.sqlite import SqlAgentStorage
from phi.knowledge.pdf import PDFKnowledgeBase
from phi.vectordb.qdrant import Qdrant
from phi.tools.duckduckgo import DuckDuckGo
from dotenv import load_dotenv
import os

# Load environment variables
load_dotenv()

# FastAPI app initialization
app = FastAPI()

# Dummy data for agents
data = {
    "agent_a": {
        "pdf_path": "dummy_path_agent_a.pdf",
        "role": "You are Agent A, a general-purpose assistant with expertise in dummy topics.",
        "description": "Agent A is designed to handle dummy queries related to various dummy topics.",
        "tools": [],
        "instructions": [
            "Always search your knowledge base before seeking external information.",
            "If the information is not found, search the web.",
            "Provide clear and concise answers."
        ]
    },
    "agent_b": {
        "pdf_path": "dummy_path_agent_b.pdf",
        "role": "You are Agent B, a specialized assistant for dummy operations and tasks.",
        "description": "Agent B specializes in providing guidance for dummy workflows and dummy operations.",
        "tools": [],
        "instructions": [
            "Search the knowledge base for answers before looking elsewhere.",
            "Ensure all responses are formatted clearly and accurately.",
            "Ask clarifying questions if the input query is unclear."
        ]
    }
}

# Initialize necessary components at startup
@app.on_event("startup")
async def startup():
    global vector_db_url, vector_db_api_key
    vector_db_url = os.getenv("VECTOR_DB_URL", "http://dummy-vector-db-url.com")
    vector_db_api_key = os.getenv("VECTOR_DB_API_KEY", "dummy_api_key")

# POST endpoint for chat
@app.post("/chat")
async def chat(
    user_name: str = Form(...),
    agent_id: str = Form(...),
    session_id: str = Form(...),
    text: str = Form(...),
):
    data_res = data.get(agent_id)
    if not data_res:
        return {"error": "Agent not found"}

    # Set up Qdrant VectorDB and Knowledge Base
    vector_db = Qdrant(
        collection=agent_id,
        url=vector_db_url,
        api_key=vector_db_api_key
    )
    knowledge_base = PDFKnowledgeBase(path=data_res["pdf_path"], vector_db=vector_db)

    # Define the main agent
    knowledge_agent = Agent(
        name=f"{agent_id} Knowledge Base",
        role=data_res["role"],
        model=OpenAIChat(id="gpt-3.5-turbo"),
        session_id=session_id,
        instructions=data_res["instructions"],
        knowledge_base=knowledge_base,
        markdown=True,
    )
    web_agent = Agent(
        name="Web Agent",
        role="Search the web for information",
        model=OpenAIChat(id="gpt-3.5-turbo"),
        tools=[DuckDuckGo()],
        instructions=["Always include sources"],
        markdown=True,
    )
    main_agent = Agent(
        name=agent_id,
        team=[knowledge_agent, web_agent],
        description=data_res["description"],
        model=OpenAIChat(id="gpt-4"),
        markdown=True,
        storage=SqlAgentStorage(table_name=f"{agent_id}_storage", db_file="agents.db"),
    )

    response = main_agent.run(text).content
    return {"response": response}