app.py

# GET /models
# GET /tools
# POST /chat -> Groq -> response

import asyncio
import json
import traceback
from typing import List, Optional
from contextlib import AsyncExitStack
import uuid

from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from fastapi import FastAPI, Request, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse


from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

from groq import Groq, APIConnectionError
from dotenv import load_dotenv
import os
import httpx

sessions = {}
unique_apikeys = []

class MCPClient:
    def __init__(self):
        self.session: Optional[ClientSession] = None
        self.exit_stack = AsyncExitStack()
        self.current_model = None
        self.groq = None
        self.api_key = None
        self.messages = [{
            "role": "system",
            "content": "You are a helpful assistant that have access to different tools via MCP. Make complete answers."
        }]
        self.tool_use = True
        self.models = None
        self.tools = []

    async def connect(self, api_key: str):
        try:
            self.groq = Groq(api_key=api_key, http_client=httpx.Client(verify=False, timeout=30))
            self.api_key = api_key
        except APIConnectionError as e:
            traceback.print_exception(e)
            return False
        except Exception as e:
            traceback.print_exception(e)
            return False
        server_params = StdioServerParameters(command="uv", args=["run", "server.py"])
        
        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
        self.stdio, self.write = stdio_transport
        self.session = await self.exit_stack.enter_async_context(ClientSession(self.stdio, self.write))
        
        await self.session.initialize()
        
        response = await self.session.list_tools()
        tools = response.tools
        print("\nConnected to server with tools:", [tool.name for tool in tools])
        self.tools = [{"type": "function", "function": { 
            "name": tool.name,
            "description": tool.description,
            "parameters": tool.inputSchema
        }} for tool in tools]

    def populate_model(self):
        self.models = sorted([m.id for m in self.groq.models.list().data])

    async def process_query(self, query: str) -> str:
        """Process a query using Groq and available tools"""
        self.messages.extend([
            {
                "role": "user",
                "content": query
            }
        ])

        response = self.groq.chat.completions.create(
            model=self.current_model,
            messages=self.messages,
            tools=self.tools,
            temperature=0
        ) if self.tool_use else self.groq.chat.completions.create(
            model=self.current_model,
            messages=self.messages,
            temperature=0.7
        )

        # Process response and handle tool calls
        final_text = []

        for choice in response.choices:
            content = choice.message.content
            tool_calls = choice.message.tool_calls
            if content:
                final_text.append(content)
            if tool_calls:
                print(tool_calls)
                for tool in tool_calls:
                    tool_name = tool.function.name
                    tool_args = tool.function.arguments

                    result = await self.session.call_tool(tool_name, json.loads(tool_args))
                    print(f"[Calling tool {tool_name} with args {tool_args}]")

                    if content is not None:
                        self.messages.append({
                            "role": "assistant",
                            "content": content
                        })
                    self.messages.append({
                        "role": "tool",
                        "tool_call_id": tool.id,
                        "content": str(result.content)
                    })

                    print(result.content[0].text)

                    response = self.groq.chat.completions.create(
                        model=self.current_model,
                        messages=self.messages,
                        temperature=0.7
                    )

                    final_text.append(response.choices[0].message.content)
        return "\n".join(final_text)

app = FastAPI()
app.add_middleware(CORSMiddleware, allow_credentials=True, allow_headers=["*"], allow_methods=["*"], allow_origins=["*"])
app.mount("/static", StaticFiles(directory="static"), name="static")
mcp = MCPClient()

class InitRequest(BaseModel):
    api_key: str

class InitResponse(BaseModel):
    success: bool
    session_id: str
    models: Optional[list] = None
    error: Optional[str] = None

class LogoutRequest(BaseModel):
    session_id: str

def get_mcp_client(session_id: str) -> MCPClient|None:
    """Get the MCPClient for a given session_id, or raise 404."""
    client = sessions.get(session_id)
    if client is None:
        raise HTTPException(status_code=404, detail="Invalid session_id. Please re-initialize.")
    return client

@app.get("/")
def root():
    return FileResponse("index.html")

@app.post("/init", response_model=InitResponse)
async def init_server(req: InitRequest):
    """
    Initializes a new MCP client session. Returns a session_id.
    """
    api_key = req.api_key
    session_id = str(uuid.uuid4())
    mcp = MCPClient()

    try:
        ok = await mcp.connect(api_key)
        if ok is False:
            raise RuntimeError("Failed to connect to MCP or Groq with API key.")
        mcp.populate_model()
        
        sessions[session_id] = mcp
        if api_key not in unique_apikeys:
            unique_apikeys.append(api_key)
        else:
            raise Exception("Session with this API key already exists. We won't re-return you the session ID. Bye-bye Hacker !!")
        return InitResponse(
            session_id=session_id,
            models=mcp.models,
            error=None,
            success=True
        )
    except Exception as e:
        traceback.print_exception(e)
        return InitResponse(
            session_id="",
            models=None,
            error=str(e),
            success=False
        )

class ChatRequest(BaseModel):
    session_id: str
    query: str
    tool_use: Optional[bool] = True
    model: Optional[str] = "llama-3.3-70b-versatile"

class ChatResponse(BaseModel):
    output: str
    error: Optional[str] = None

@app.post("/chat", response_model=ChatResponse)
async def chat(req: ChatRequest):
    """
    Handles chat requests for a given session.
    """
    try:
        mcp = get_mcp_client(req.session_id)
        mcp.tool_use = req.tool_use
        if req.model in mcp.models:
            mcp.current_model = req.model
        else:
            raise ValueError(f"Model not recognized: Not in the model list: {mcp.models}")
        result = await mcp.process_query(req.query)
        return ChatResponse(output=result)
    except Exception as e:
        traceback.print_exception(e)
        return ChatResponse(output="", error=str(e))
    
@app.post("/logout")
async def logout(logout_req: LogoutRequest):
    """Clean up session resources."""
    mcp = sessions.pop(logout_req.session_id, None)
    unique_apikeys.remove(mcp.api_key)
    if mcp and hasattr(mcp.exit_stack, "aclose"):
        await mcp.exit_stack.aclose()
    return {"success": True}