main.py

import os
import json
from typing import List, Optional, Any
from fastapi import FastAPI, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, StreamingResponse
from pydantic import BaseModel
import anthropic
import uvicorn
from tools import TOOLS, ToolExecutor

# Initialize FastAPI app
app = FastAPI(title="Data Engineering Agent")

# Models for API
class Message(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    messages: List[Message]
    stream: bool = False

class ChatResponse(BaseModel):
    message: Message

# Get API key from environment
ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
if not ANTHROPIC_API_KEY:
    print("Warning: ANTHROPIC_API_KEY not set. Chat functionality will not work.")

# Initialize Anthropic client
client = anthropic.Anthropic(api_key=ANTHROPIC_API_KEY) if ANTHROPIC_API_KEY else None

# Load system prompt from file
def load_system_prompt() -> str:
    """Load the system prompt from system_prompt.md"""
    prompt_path = os.path.join(os.path.dirname(__file__), "system_prompt.md")
    with open(prompt_path, "r") as f:
        return f.read().strip()

SYSTEM_PROMPT = load_system_prompt()

@app.get("/", response_class=HTMLResponse)
async def root():
    """Serve the chat interface"""
    with open("static/index.html", "r") as f:
        return f.read()

@app.get("/health")
async def health():
    """Health check endpoint"""
    return {"status": "healthy"}

@app.post("/api/chat")
async def chat(request: ChatRequest):
    """Handle chat requests with tool support"""
    if not client:
        raise HTTPException(status_code=500, detail="API key not configured")

    try:
        # Convert messages to Anthropic format
        anthropic_messages = [
            {"role": msg.role, "content": msg.content}
            for msg in request.messages
        ]

        if request.stream:
            # Streaming response with tool use
            async def generate():
                current_messages = anthropic_messages.copy()

                while True:
                    # Call Claude with tools
                    response = client.messages.create(
                        model="claude-sonnet-4-5-20250929",
                        max_tokens=8192,
                        system=SYSTEM_PROMPT,
                        messages=current_messages,
                        tools=TOOLS
                    )

                    # Process response content
                    for block in response.content:
                        if block.type == "text":
                            # Stream text content
                            yield f"data: {json.dumps({'content': block.text})}\n\n"

                        elif block.type == "tool_use":
                            # Notify that tool is being used
                            tool_use_msg = {
                                'tool_use': {
                                    'name': block.name,
                                    'input': block.input
                                }
                            }
                            yield f"data: {json.dumps(tool_use_msg)}\n\n"

                            # Execute the tool
                            tool_result = ToolExecutor.execute(block.name, block.input)

                            # Send tool result status
                            tool_result_msg = {
                                'tool_result': {
                                    'name': block.name,
                                    'success': tool_result.get('success', False)
                                }
                            }
                            yield f"data: {json.dumps(tool_result_msg)}\n\n"

                    # Check if we need to continue (if there were tool uses)
                    tool_uses = [block for block in response.content if block.type == "tool_use"]

                    if not tool_uses:
                        # No tool uses, we're done
                        break

                    # Add assistant response with tool uses to messages
                    current_messages.append({
                        "role": "assistant",
                        "content": response.content
                    })

                    # Add tool results to messages
                    tool_results = []
                    for block in response.content:
                        if block.type == "tool_use":
                            result = ToolExecutor.execute(block.name, block.input)
                            tool_results.append({
                                "type": "tool_result",
                                "tool_use_id": block.id,
                                "content": json.dumps(result)
                            })

                    current_messages.append({
                        "role": "user",
                        "content": tool_results
                    })

                # Send completion signal
                yield f"data: {json.dumps({'done': True})}\n\n"

            return StreamingResponse(
                generate(),
                media_type="text/event-stream"
            )
        else:
            # Non-streaming response with tool use
            current_messages = anthropic_messages.copy()

            while True:
                response = client.messages.create(
                    model="claude-sonnet-4-5-20250929",
                    max_tokens=8192,
                    system=SYSTEM_PROMPT,
                    messages=current_messages,
                    tools=TOOLS
                )

                # Check if there are tool uses
                tool_uses = [block for block in response.content if block.type == "tool_use"]

                if not tool_uses:
                    # No tool uses, return the text response
                    text_content = " ".join([
                        block.text for block in response.content
                        if block.type == "text"
                    ])
                    return ChatResponse(
                        message=Message(
                            role="assistant",
                            content=text_content
                        )
                    )

                # Execute tools and continue conversation
                current_messages.append({
                    "role": "assistant",
                    "content": response.content
                })

                tool_results = []
                for block in response.content:
                    if block.type == "tool_use":
                        result = ToolExecutor.execute(block.name, block.input)
                        tool_results.append({
                            "type": "tool_result",
                            "tool_use_id": block.id,
                            "content": json.dumps(result)
                        })

                current_messages.append({
                    "role": "user",
                    "content": tool_results
                })

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Mount static files
app.mount("/static", StaticFiles(directory="static"), name="static")

def main():
    port = int(os.getenv("PORT", "50051"))
    hostname = os.getenv("TOWER__HOSTNAME", "(local)")
    print(f"Starting Data Engineering Agent on {hostname}")
    print(f"Server running on http://0.0.0.0:{port}")

    try:
        uvicorn.run(
            app,
            host="0.0.0.0",
            port=port,
            log_level="info",
            use_colors=False
        )
    except Exception as e:
        print(f"Failed to start server: {str(e)}")
        raise
    finally:
        print("FastAPI application shutting down...")

if __name__ == "__main__":
    main()