main.py

import asyncio
import uuid
from typing import Optional, Dict, List, Literal
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from enum import Enum
from agent import Agent
from pathlib import Path
from dotenv import load_dotenv
import os
import sys
import openai
from utcp.client.utcp_client_config import UtcpClientConfig, UtcpDotEnv
from utcp.client.utcp_client import UtcpClient
from vector_store import EmbeddingInMemRepo
from embedding_model import EmbeddingModel
from tool_selector import ToolSelector
from tag_analyzer import TagAnalyzer
from utcp.shared.utcp_manual import UtcpManual
from utcp.shared.tool import utcp_tool
from utcp.shared.provider import HttpProvider
import aiohttp
from utcp.shared.tool import ToolInputOutputSchema
from pydantic.json_schema import GenerateJsonSchema

load_dotenv(Path(__file__).parent / ".env")

if not os.environ.get("OPENAI_API_KEY"):
    print("Error: OPENAI_API_KEY not found in environment variables")
    print("Please set it in the .env file")
    sys.exit(1)

providers_file_path = str(Path(__file__).parent / "providers.json")

api_key = os.environ.get("OPENAI_API_KEY")

# Create a configuration for the UTCP client
config = UtcpClientConfig(
    providers_file_path=providers_file_path,
    load_variables_from=[
        UtcpDotEnv(env_file_path=str(Path(__file__).parent / ".env"))
    ]
)
openai_client = openai.AsyncOpenAI(api_key=api_key)

embedding_model = EmbeddingModel(api_key=api_key)
tool_repo = EmbeddingInMemRepo(embedding_model)
tool_search_strategy = ToolSelector(tool_repo, embedding_model, TagAnalyzer(openai_client=openai_client))
utcp_client = None
agent = None

# FastAPI app setup
__version__ = "1.0.0"
BASE_PATH = "http://localhost:1646"

app = FastAPI(title="UTCP Agent API", version=__version__)

class RegisterProviderResponse(BaseModel):
    success: bool
    message: str
    provider_name: str

generator = GenerateJsonSchema()
resolved_schema = generator.generate(HttpProvider.__pydantic_core_schema__)
resolved_schema = generator.resolve_ref_schema(resolved_schema)
resolved_schema["required"] = ["name", "url"]
tool_schema = ToolInputOutputSchema(
    properties={"provider": resolved_schema},
    required=["provider"],
    title="Register Provider"
)

@utcp_tool(
    tool_provider=HttpProvider(
        name="utcp_agent",
        http_method="POST",
        url="http://localhost:1646/register-provider",
        body_field="provider"
    ),
    tags=["Register", "Tools", "New Provider", "HTTP", "Provider", "Manual", "OpenAPI", "UTCP", "UTCP Agent"],
    description="Register a new HTTP tool provider for yourself",
    inputs=tool_schema
)
@app.post("/register-provider", response_model=RegisterProviderResponse)
async def register_http_provider(provider: HttpProvider) -> RegisterProviderResponse:
    """Register an HTTP tool provider with the UTCP client"""
    global utcp_client
    
    if utcp_client is None:
        raise HTTPException(status_code=500, detail="UTCP client not initialized")
    
    try:
        # Register the provider with the UTCP client
        await utcp_client.register_tool_provider(provider)
        
        return RegisterProviderResponse(
            success=True,
            message=f"HTTP provider '{provider.name}' registered successfully",
            provider_name=provider.name
        )
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Failed to register provider: {str(e)}"
        )

# API endpoints
@app.get("/utcp", response_model=UtcpManual)
def get_utcp():
    """Get UTCP manual information"""
    manual = UtcpManual.create(version=__version__)
    return manual

@app.get("/health")
def health_check():
    """Health check endpoint"""
    return {"status": "healthy", "version": __version__}

async def wait_for_server_ready(max_attempts=30, delay=1):
    """Wait for the FastAPI server to be ready"""
    for attempt in range(max_attempts):
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(f"{BASE_PATH}/health") as response:
                    if response.status == 200:
                        return True
        except Exception:
            pass
        await asyncio.sleep(delay)
    return False

async def run_interactive_agent():
    """Run the interactive agent in a separate task"""
    global utcp_client, agent
    
    # Wait for server to be ready
    print("Waiting for server to start...")
    if await wait_for_server_ready():
        print("Server is ready!")
    else:
        print("Warning: Server may not be fully ready")
    
    # Initialize UTCP client first
    utcp_client = await UtcpClient.create(config, tool_repository=tool_repo, search_strategy=tool_search_strategy)
    agent = Agent(utcp_client, openai_client)
    
    print("\n=== Interactive Agent Started ===")
    print("The agent now has access to itself via the API at http://localhost:1646")
    print("You can register HTTP providers via POST /register-provider")
    print("Available endpoints:")
    print("  - GET /utcp (UTCP manual)")
    print("  - POST /register-provider (Register HTTP provider)")
    print("  - GET /health (Health check)")
    print("Type 'exit' or 'quit' to stop\n")
    
    while True:
        try:
            user_prompt = await asyncio.get_event_loop().run_in_executor(None, input, "User: ")
            
            if user_prompt.lower() in ["exit", "quit"]:
                print("Shutting down...")
                break

            await agent.chat(user_prompt)
        except KeyboardInterrupt:
            print("\nShutting down...")
            break
        except Exception as e:
            print(f"Error in agent chat: {e}")

async def run_server():
    """Run the FastAPI server"""
    import uvicorn
    config = uvicorn.Config(app, host="localhost", port=1646, log_level="info")
    server = uvicorn.Server(config)
    await server.serve()

async def main():
    """Main function to run both FastAPI server and interactive agent"""

    # Run both server and interactive agent concurrently
    await asyncio.gather(
        run_server(),
        run_interactive_agent()
    )

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nApplication stopped by user")