from fastapi import FastAPI
from pydantic import BaseModel
from phi.agent import Agent
from phi.model.groq import Groq
from phi.tools.yfinance import YFinanceTools
from phi.tools.duckduckgo import DuckDuckGo
import os
from dotenv import load_dotenv

import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Load environment variables
load_dotenv()
gorq_api_key = os.getenv("GROQ_API_KEY")

# Initialize FastAPI app
app = FastAPI()

# Define input schema
class QueryRequest(BaseModel):
    input_text: str

# Web Search Agent
web_search_agent = Agent(
    name="Web Search Agent",
    role="Search the web for the information",
    model=Groq(id="llama-3.3-70b-versatile",api_key=gorq_api_key),
    tools=[DuckDuckGo()],
    instructions=["Always include sources"],
    show_tools_calls=True,
    markdown=True,
)

# Financial Agent
finance_agent = Agent(
    name="Finance AI Agent",
    model=Groq(
        id="llama-3.3-70b-versatile",
        api_key=gorq_api_key,
        role="Get financial data",
    ),
    tools=[
        YFinanceTools(
            stock_price=True,
            analyst_recommendations=True,
            stock_fundamentals=True,
            company_news=True,
        ),
    ],
    instructions=["Use tables to display the data"],
    show_tool_calls=True,
    markdown=True,
)

# Multi-Agent
multi_ai_agent = Agent(
    model=Groq(id="llama-3.3-70b-versatile", api_key=gorq_api_key),
    team=[web_search_agent, finance_agent],
    instructions=["Always include sources", "Use tables to display the data"],
    show_tool_calls=True,
    markdown=True,
)

@app.post("/query")
async def process_query(request: QueryRequest):
    try:
        response = multi_ai_agent.run(request.input_text)
        return {"response": response}
    except Exception as e:
        logger.error(f"Error processing query: {e}")
        return {"error": f"An error occurred: {str(e)}"}