server1.py

import os
import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from dotenv import load_dotenv

from langchain_groq import ChatGroq
from langchain_community.tools.tavily_search import TavilySearchResults
from langgraph.prebuilt import create_react_agent
from langchain_core.messages import HumanMessage, SystemMessage

load_dotenv()

app = FastAPI(title="A2A Agent Server")

# Model configuration
llm = ChatGroq(model="llama-3.3-70b-versatile", temperature=0)

# Tools
tavily_tool = TavilySearchResults(max_results=3)

# System Prompts
SYS_PROMPT_1 = "You are a specialized agent providing details about Microsoft products. Use the search tool to find accurate and up-to-date information about Microsoft products. Answer ONLY questions related to Microsoft products."
SYS_PROMPT_2 = "You are a specialized agent providing support, clarification, and formulating FAQs for Microsoft. Use the search tool to find accurate and up-to-date support information. Answer ONLY questions related to Microsoft support and FAQs."
SYS_PROMPT_3 = "You are a strict fallback agent. Your only job is to firmly reject the user's request. State clearly that you can only provide responses for queries related to Microsoft products, support, or services. Do NOT answer any other questions."

# Agents with proper system messages (using prompt instead of state_modifier)
agent_1 = create_react_agent(llm, tools=[tavily_tool], prompt=SYS_PROMPT_1)
agent_2 = create_react_agent(llm, tools=[tavily_tool], prompt=SYS_PROMPT_2)
agent_3 = create_react_agent(llm, tools=[], prompt=SYS_PROMPT_3)  # No tools needed for rejection

agents_map = {
    "agent_1": agent_1,
    "agent_2": agent_2,
    "agent_3": agent_3,
}

# A2A Agent Card (Discovery)
AGENT_CARD = {
    "name": "Microsoft Agentic System",
    "description": "A system of agents specialized in Microsoft products and support.",
    "agents": [
        {
            "id": "agent_1",
            "name": "Product Agent",
            "description": "Details about Microsoft products."
        },
        {
            "id": "agent_2",
            "name": "Support Agent",
            "description": "Microsoft support, clarifications, and FAQs."
        },
        {
            "id": "agent_3",
            "name": "Reject Agent",
            "description": "Fallback agent for non-Microsoft queries."
        }
    ]
}

class QueryRequest(BaseModel):
    query: str
    session_id: str

@app.get("/.well-known/agent.json")
def get_agent_card():
    return AGENT_CARD

@app.post("/agent/{agent_id}")
async def run_agent(agent_id: str, request: QueryRequest):
    if agent_id not in agents_map:
        raise HTTPException(status_code=404, detail="Agent not found")
    
    agent = agents_map[agent_id]
    
    # Run the agent
    response = await agent.ainvoke({"messages": [HumanMessage(content=request.query)]})
    
    # Get the last message content
    result_text = response["messages"][-1].content
    
    return {"agent_id": agent_id, "response": result_text}

if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=8001)