from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import requests
import json
from typing import List, Optional
from langchain.tools import BaseTool
from langchain.agents import initialize_agent
from langchain.llms import Ollama
from langchain.memory import ConversationBufferMemory
from langchain.agents import AgentType

class Tool(BaseModel):
    name: str
    description: str
    function: str

class GenerationRequest(BaseModel):
    prompt: str
    max_tokens: int = 100
    temperature: float = 0.7
    tools: Optional[List[Tool]] = None

app = FastAPI()

OLLAMA_URL = "http://localhost:11434/api/generate"

@app.post("/generate")
async def generate_code(request: GenerationRequest):
    try:
        # Initialize Ollama LLM
        llm = Ollama(
            base_url="http://localhost:11434",
            model="codellama",
            temperature=request.temperature,
            max_tokens=request.max_tokens
        )

        # Initialize memory
        memory = ConversationBufferMemory(memory_key="chat_history")

        # Initialize tools if provided
        tools = []
        if request.tools:
            for tool in request.tools:
                tools.append(
                    BaseTool(
                        name=tool.name,
                        description=tool.description,
                        func=lambda x: eval(tool.function)(x)
                    )
                )

        # Initialize agent
        agent = initialize_agent(
            tools=tools,
            llm=llm,
            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
            memory=memory
        )

        # Generate response
        response = agent.run(request.prompt)
        
        return {"generated_code": response}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)