refactor(main): improve code structure and error handling

Refactored the main application file to enhance code readability and maintainability. Introduced logging for better error tracking and replaced JSONResponse with HTTPException for more consistent error handling. Added detailed docstrings to functions and methods for better understanding of their purpose and usage. Also, restructured the code to separate concerns and improve the logical flow.

Additionally, created a new `schemas.py` file to define data models using Pydantic, which helps in validating and organizing request and response data structures.

This refactor aims to improve the overall robustness and scalability of the application by ensuring that the code is more modular and easier to maintain.

main.py

@@ -1,26 +1,52 @@
-import os
-from fastapi import FastAPI, HTTPException, Depends, Header, Request, Response
-from fastapi.responses import JSONResponse, StreamingResponse
-from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
-from pydantic import BaseModel
-from typing import List, Optional
-from anthropic import Anthropic
 import json
+import logging
+import os
 import time
+from concurrent.futures import ThreadPoolExecutor
 from contextvars import ContextVar
+from typing import Any, Dict, Generator, List
 
+from anthropic import Anthropic
+from fastapi import FastAPI, HTTPException, Request, Response
+from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi.security import HTTPBearer
+from starlette.concurrency import run_in_threadpool
+
+from schemas import OpenAIChatCompletionForm, FilterForm
+
+# logger
+logger = logging.getLogger()
+
+# FastAPI app initialization
 app = FastAPI()
 security = HTTPBearer()
 
-# Context variable to store the token
+# Context variable for token storage
 token_context = ContextVar('token', default=None)
 
-# Liste des endpoints qui ne nécessitent pas d'authentification
+# Endpoints that don't require authentication
 PUBLIC_ENDPOINTS = {"/"}
 
+# Available Anthropic models
+AVAILABLE_MODELS = [
+    "claude-3-haiku-20240307",
+    "claude-3-opus-20240229",
+    "claude-3-sonnet-20240229",
+    "claude-3-5-sonnet-20241022"
+]
+
 @app.middleware("http")
 async def auth_middleware(request: Request, call_next):
-    # Skip authentication for public endpoints
+    """
+    Middleware for handling authentication and response logging.
+    
+    Args:
+        request: The incoming HTTP request
+        call_next: The next middleware in the chain
+        
+    Returns:
+        Response: The processed HTTP response
+    """
     if request.url.path in PUBLIC_ENDPOINTS:
         start_time = time.perf_counter()
         response = await call_next(request)
@@ -31,22 +57,20 @@ async def auth_middleware(request: Request, call_next):
     try:
         auth_header = request.headers.get('Authorization')
         if not auth_header:
-            return JSONResponse(
+            raise HTTPException(
                 status_code=401,
-                content={"detail": "No authorization header"}
+                detail="No authorization header"
             )
         
         scheme, token = auth_header.split()
         if scheme.lower() != 'bearer':
-            return JSONResponse(
+            raise HTTPException(
                 status_code=401,
-                content={"detail": "Invalid authentication scheme"}
+                detail="Invalid authentication scheme"
             )
         
-        # Store token in context
         token_context.set(token)
         
-        # Add processing time header
         start_time = time.perf_counter()
         response = await call_next(request)
         process_time = time.perf_counter() - start_time
@@ -54,46 +78,60 @@ async def auth_middleware(request: Request, call_next):
         
         return response
             
+    except HTTPException as http_ex:
+        logger.error(
+            f"HTTP Exception - Status: {http_ex.status_code} - "
+            f"Detail: {http_ex.detail} - Path: {request.url.path}"
+        )
+        return JSONResponse(
+            status_code=http_ex.status_code,
+            content={"detail": http_ex.detail}
+        )
     except Exception as e:
+        logger.error(
+            f"Unexpected error in middleware - Error: {str(e)} - "
+            f"Path: {request.url.path}",
+            exc_info=True
+        )
         return JSONResponse(
-            status_code=401,
-            content={"detail": "Invalid authorization header"}
+            status_code=500,
+            content={"detail": "Internal server error"}
         )
 
-# Function to get Anthropic client with current token
+
 def get_anthropic_client():
+    """
+    Get an authenticated Anthropic client using the current token.
+    
+    Returns:
+        Anthropic: An authenticated Anthropic client instance
+        
+    Raises:
+        HTTPException: If no authorization token is found
+    """
     token = token_context.get()
     if not token:
         raise HTTPException(status_code=401, detail="No authorization token found")
     return Anthropic(api_key=token)
 
-# Available models
-AVAILABLE_MODELS = [
-    "claude-3-haiku-20240307",
-    "claude-3-opus-20240229",
-    "claude-3-sonnet-20240229",
-    "claude-3-5-sonnet-20241022"
-]
-
-class Message(BaseModel):
-    role: str
-    content: str
-
-class ChatCompletionRequest(BaseModel):
-    model: str
-    messages: List[Message]
-    stream: bool = False
-    max_tokens: Optional[int] = 1024
 
+@app.get("/v1")
 @app.get("/")
 async def read_root():
+    """Root endpoint for API health check."""
     return {"Hello": "World!"}
 
+
+@app.get("/v1/models")
 @app.get("/models")
 async def get_models():
-    """Get available Anthropic models."""
-    # Test the token by creating a client
-    get_anthropic_client()
+    """
+    Get available Anthropic models.
+    
+    Returns:
+        JSONResponse: List of available models and their details
+    """
+    get_anthropic_client()  # Verify token validity
     
     models = [
         {
@@ -115,97 +153,178 @@ async def get_models():
         }
     )
 
-@app.post("/v1/chat/completions")
-async def create_chat_completion(request: ChatCompletionRequest):
-    """Generate chat completions using Anthropic models."""
-    try:
-        if request.stream:
-            return StreamingResponse(
-                stream_response(request),
-                media_type="text/event-stream"
-            )
-        else:
-            return await generate_completion(request)
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
 
-async def generate_completion(request: ChatCompletionRequest):
-    """Generate a non-streaming completion."""
-    messages = [{"role": m.role, "content": m.content} for m in request.messages]
+def stream_message(
+    model: str,
+    messages: List[Dict[str, Any]]
+) -> Generator[str, None, None]:
+    """
+    Stream messages using the specified model.
     
-    # Get client with current token
+    Args:
+        model: The model identifier to use
+        messages: List of messages to process
+        
+    Returns:
+        Generator: Stream of SSE formatted responses
+    """
     client = get_anthropic_client()
-    
     response = client.messages.create(
-        model=request.model,
-        max_tokens=request.max_tokens,
+        model=model,
+        max_tokens=1024,
+        messages=messages,
+        stream=True
+    )
+
+    def event_stream() -> Generator[str, None, None]:
+        message_id = None
+        
+        for chunk in response:
+            if not message_id:
+                message_id = f"chatcmpl-{int(time.time())}"
+                
+            if chunk.type == 'content_block_delta':
+                data = {
+                    "id": message_id,
+                    "object": "chat.completion.chunk",
+                    "created": int(time.time()),
+                    "model": model,
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {
+                                "content": (
+                                    chunk.delta.text
+                                    if hasattr(chunk.delta, 'text')
+                                    else ""
+                                )
+                            },
+                            "logprobs": None,
+                            "finish_reason": None,
+                        }
+                    ],
+                }
+                yield f"data: {json.dumps(data)}\n\n"
+
+            elif chunk.type == 'content_block_stop':
+                data = {
+                    "id": message_id,
+                    "object": "chat.completion.chunk",
+                    "created": int(time.time()),
+                    "model": model,
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {},
+                            "logprobs": None,
+                            "finish_reason": "stop",
+                        }
+                    ],
+                }
+                yield f"data: {json.dumps(data)}\n\n"
+
+        yield "data: [DONE]\n\n"
+
+    return event_stream()
+
+
+def send_message(model: str, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
+    """
+    Send a message via the Anthropic provider without streaming.
+
+    Args:
+        model: The model identifier to use
+        messages: List of messages to process
+
+    Returns:
+        dict: The formatted response from the model
+    """
+    client = get_anthropic_client()
+    response = client.messages.create(
+        model=model,
+        max_tokens=1024,
         messages=messages
     )
 
+    content = response.content[0].text if response.content else ""
+
     return {
         "id": response.id,
         "object": "chat.completion",
         "created": int(time.time()),
-        "model": request.model,
-        "choices": [{
-            "index": 0,
-            "message": {
-                "role": "assistant",
-                "content": response.content[0].text if response.content else "",
-            },
-            "finish_reason": "stop"
-        }],
-        "usage": {
-            "prompt_tokens": response.usage.input_tokens,
-            "completion_tokens": response.usage.output_tokens,
-            "total_tokens": response.usage.input_tokens + response.usage.output_tokens
-        }
+        "model": model,
+        "choices": [
+            {
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": content,
+                },
+                "logprobs": None,
+                "finish_reason": "stop",
+            }
+        ],
     }
 
-async def stream_response(request: ChatCompletionRequest):
-    """Stream the completion response."""
-    messages = [{"role": m.role, "content": m.content} for m in request.messages]
+
+@app.post("/v1/chat/completions")
+@app.post("/chat/completions")
+async def generate_chat_completion(form_data: OpenAIChatCompletionForm):
+    """
+    Generate chat completions from the model.
     
-    # Get client with current token
-    client = get_anthropic_client()
+    Args:
+        form_data: The chat completion request parameters
+        
+    Returns:
+        Union[StreamingResponse, dict]: Either a streaming response or a complete message
+    """
+    messages = [
+        {"role": message.role, "content": message.content}
+        for message in form_data.messages
+    ]
+    model = form_data.model
+
+    def job():
+        """Handle both streaming and non-streaming modes."""
+        if form_data.stream:
+            return StreamingResponse(
+                stream_message(model=model, messages=messages),
+                media_type="text/event-stream"
+            )
+        return send_message(model=model, messages=messages)
+
+    with ThreadPoolExecutor() as executor:
+        return await run_in_threadpool(job)
+
+
+@app.post("/v1/{pipeline_id}/filter/inlet")
+@app.post("/{pipeline_id}/filter/inlet")
+async def filter_inlet(pipeline_id: str, form_data: FilterForm):
+    """
+    Handle inlet filtering for the pipeline.
     
-    response = client.messages.create(
-        model=request.model,
-        max_tokens=request.max_tokens,
-        messages=messages,
-        stream=True
-    )
+    Args:
+        pipeline_id: The ID of the pipeline
+        form_data: The filter parameters
+        
+    Returns:
+        dict: The processed request body
+    """
+    return form_data.body
 
-    for chunk in response:
-        if chunk.type == "message_start":
-            continue
-            
-        if chunk.type == "content_block_delta":
-            data = {
-                "id": chunk.message.id,
-                "object": "chat.completion.chunk",
-                "created": int(time.time()),
-                "model": request.model,
-                "choices": [{
-                    "index": 0,
-                    "delta": {"content": chunk.delta.text if hasattr(chunk.delta, "text") else ""},
-                    "finish_reason": None
-                }]
-            }
-            yield f"data: {json.dumps(data)}\n\n"
-            
-        elif chunk.type == "content_block_stop":
-            data = {
-                "id": chunk.message.id,
-                "object": "chat.completion.chunk",
-                "created": int(time.time()),
-                "model": request.model,
-                "choices": [{
-                    "index": 0,
-                    "delta": {},
-                    "finish_reason": "stop"
-                }]
-            }
-            yield f"data: {json.dumps(data)}\n\n"
+
+@app.post("/v1/{pipeline_id}/filter/outlet")
+@app.post("/{pipeline_id}/filter/outlet")
+async def filter_outlet(pipeline_id: str, form_data: FilterForm):
+    """
+    Handle outlet filtering for the pipeline.
     
-    yield "data: [DONE]\n\n"
+    Args:
+        pipeline_id: The ID of the pipeline
+        form_data: The filter parameters
+        
+    Returns:
+        dict: The processed request body
+    """
+    return form_data.body

schemas.py

@@ -0,0 +1,41 @@
+from typing import List, Optional, Union
+from pydantic import BaseModel, ConfigDict, RootModel
+
+
+class ImageContent(BaseModel):
+    """Model for image content in messages."""
+    type: str
+    image_url: dict
+
+
+class TextContent(BaseModel):
+    """Model for text content in messages."""
+    type: str
+    text: str
+
+
+class MessageContent(RootModel):
+    """Model for message content that can be either text or image."""
+    root: Union[TextContent, ImageContent]
+
+
+class OpenAIChatMessage(BaseModel):
+    """Model for chat messages in OpenAI format."""
+    role: str
+    content: Union[str, List[MessageContent]]
+    model_config = ConfigDict(extra="allow")
+
+
+class OpenAIChatCompletionForm(BaseModel):
+    """Model for chat completion request parameters."""
+    stream: bool = True
+    model: str
+    messages: List[OpenAIChatMessage]
+    model_config = ConfigDict(extra="allow")
+
+
+class FilterForm(BaseModel):
+    """Model for filter request parameters."""
+    body: dict
+    user: Optional[dict] = None
+    model_config = ConfigDict(extra="allow")