#!/usr/bin/env python3.12 import importlib import logging import regex from aiohttp import ClientSession, ClientTimeout from fastapi import FastAPI, Security, Request, HTTPException from fastapi.security import APIKeyHeader, APIKeyQuery from pydantic import BaseModel class AI(FastAPI): """AI Endpoints""" def __init__(self, app: FastAPI, my_util, constants, glob_state): # pylint: disable=super-init-not-called self.app = app self.util = my_util self.constants = constants self.glob_state = glob_state self.url_clean_regex = regex.compile(r'^\/ai\/') self.endpoints = { "ai": self.ai_handler, #tbd } for endpoint, handler in self.endpoints.items(): app.add_api_route(f"/{endpoint}/{{any:path}}", handler, methods=["POST"]) async def ai_handler(self, request: Request): """ /ai/ AI Request """ if not self.util.check_key(request.url.path, request.headers.get('X-Authd-With')): raise HTTPException(status_code=403, detail="Unauthorized") """ TODO: Implement Claude Currently only routes to local LLM """ local_llm_headers = { 'Authorization': f'Bearer {self.constants.LOCAL_LLM_KEY}' } forward_path = self.url_clean_regex.sub('', request.url.path) try: async with ClientSession() as session: async with await session.post(f'{self.constants.LOCAL_LLM_HOST}/{forward_path}', json=await request.json(), headers=local_llm_headers, timeout=ClientTimeout(connect=15, sock_read=30)) as request: await self.glob_state.increment_counter('ai_requests') response = await request.json() return response except Exception as e: # pylint: disable=broad-exception-caught logging.error("Error: %s", e) return { 'err': True, 'errorText': 'General Failure' }