Create app.py

app.py

@@ -0,0 +1,332 @@
+# Models: https://github.com/abacusai/api-python/blob/main/abacusai/api_class/enums.py
+model_mapping = {
+    "sonnet": "CLAUDE_V3_5_SONNET",
+    "4o": "OPENAI_GPT4O",
+    "32": "OPENAI_GPT4_32K",
+    "turbo": "OPENAI_GPT4_128K_LATEST",
+    "vision": "OPENAI_GPT4_VISION",
+    "3.5": "OPENAI_GPT3_5",
+    "opus": "CLAUDE_V3_OPUS",
+    "haiku": "CLAUDE_V3_HAIKU",
+    "claude-2": "CLAUDE_V2_1",
+    "pro": "GEMINI_1_5_PRO",
+    "palm": "PALM",
+    "llama": "LLAMA3_LARGE_CHAT",
+    "_legacy_sonnet": "CLAUDE_V3_SONNET",
+    "_legacy_gemini": "GEMINI_PRO",
+    "_legacy_palm": "PALM_TEXT"
+}
+ 
+# requirements: fastapi, curl_cffi, cachetools, websockets, orjson, uvicorn, uvloop, slowapi
+import os
+set_env = lambda var_name, default=None: environment_variables.update({var_name: os.getenv(var_name, default)}) or os.getenv(var_name, default)
+environment_variables = {}
+ 
+# Define your environment variables using the set_env function
+FALLBACK_MODEL = set_env("FALLBACK_LLM", "CLAUDE_V3_5_SONNET")
+RATE_LIMIT = set_env("RATE_LIMIT", "1/4 second")
+LOG_LEVEL = set_env("LOG_LEVEL", "INFO")
+PORT = int(set_env("PORT", "8000"))
+BASE_HOST = set_env("BASE_HOST", "apps.abacus.ai")
+ 
+DEPLOYMENT_CACHE_TTL = 3600 * 24  # 24 hours
+IMPERSONATE_BASE = "chrome"
+CURL_MAX_CLIENTS = 300
+ 
+import asyncio
+import json
+import uuid
+import random
+import logging
+from typing import Dict, Any
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import StreamingResponse
+from curl_cffi import requests, CurlOpt, CurlHttpVersion
+ 
+from cachetools import TTLCache
+deployment_cache = TTLCache(maxsize=300, ttl=DEPLOYMENT_CACHE_TTL)
+cache_lock = asyncio.Lock()
+ 
+import websockets
+ 
+try:
+    import orjson as json
+    jsonDumps = lambda text: json.dumps(text).decode('utf-8')
+except ImportError:
+    import json
+    jsonDumps = json.dumps
+from slowapi import Limiter, _rate_limit_exceeded_handler
+from slowapi.util import get_remote_address
+from slowapi.errors import RateLimitExceeded
+ 
+CURL_OPTS = {
+    CurlOpt.TCP_NODELAY: 1, CurlOpt.FORBID_REUSE: 0, CurlOpt.FRESH_CONNECT: 0, CurlOpt.TCP_KEEPALIVE: 1, CurlOpt.MAXAGE_CONN: 30
+}
+client = requests.AsyncSession(
+    impersonate=IMPERSONATE_BASE, default_headers=True, max_clients=CURL_MAX_CLIENTS, curl_options=CURL_OPTS, http_version=CurlHttpVersion.V2_PRIOR_KNOWLEDGE
+)
+ 
+from rich.logging import RichHandler
+from rich.console import Console
+from rich.table import Table
+ 
+# Setup logger with RichHandler for better logging output
+logging.basicConfig(
+    level=getattr(logging, LOG_LEVEL),
+    format="%(message)s",
+    datefmt="[%X]",
+    handlers=[RichHandler()]
+)
+logger = logging.getLogger(__name__)
+ 
+app = FastAPI()
+ 
+limiter = Limiter(key_func=get_remote_address)
+app.state.limiter = limiter
+app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
+ 
+def convert_unicode_escape(s):
+    return s.encode('utf-8').decode('unicode-escape')
+ 
+async def make_request(method: str, url: str, headers: dict, data: dict):
+    try:
+        response = await client.request(method=method, url=url, headers=headers, json=data)
+        status = response.status_code
+        if status == 200:
+            return response
+        elif status in (401, 403):
+            raise HTTPException(status_code=401, detail="Invalid authorization info")
+        else:
+            raise HTTPException(status_code=status, detail=f"Network issue: {response.text}")
+    except Exception as e:
+        logger.error(f"Request error: {str(e)}", exc_info=True)
+        raise HTTPException(status_code=500, detail=f"Request error: {str(e)}")
+ 
+def map_model(requestModel):
+    model = requestModel.lower()
+ 
+    if model.startswith('adv/'):
+        model = model[4:]
+        return model if model else FALLBACK_MODEL
+ 
+    return next((value for key, value in model_mapping.items() if key in model), FALLBACK_MODEL)
+ 
+async def get_deployment_details(apikey: str) -> str:
+    if apikey in deployment_cache:
+        return deployment_cache[apikey]
+ 
+    async with cache_lock:
+        if apikey in deployment_cache:
+            return deployment_cache[apikey]
+ 
+        headers = {
+            'apiKey': apikey,
+            'accept': '*/*',
+        }
+ 
+        response = await make_request(
+            method="GET",
+            url=f"https://{BASE_HOST}/api/listExternalApplications",
+            headers=headers,
+            data={}
+        )
+ 
+        result = response.json()
+        logger.debug(f"List external applications result: {result}")
+ 
+        if result.get("success") and result.get("result"):
+            deployment_details = result["result"][0]
+            deployment_cache[apikey] = deployment_details
+            logger.info(f"#{deployment_details['deploymentId']} - Access granted successfully")
+            return deployment_details
+        else:
+            raise HTTPException(status_code=500, detail="Failed to retrieve deployment info")
+ 
+async def create_conversation(apikey: str) -> str:
+    deployment_details = await get_deployment_details(apikey)
+ 
+    payload = {
+        "deploymentId": deployment_details["deploymentId"],
+        "name": "New Chat",
+        "externalApplicationId": deployment_details["externalApplicationId"]
+    }
+    try:
+        headers = {
+            'Content-Type': 'application/json',
+            'apiKey': apikey,
+            'REAI-UI': '1',
+            'X-Abacus-Org-Host': 'apps'
+        }
+        response = await make_request(
+            method="POST",
+            url=f"https://{BASE_HOST}/api/createDeploymentConversation",
+            headers=headers,
+            data=payload
+        )
+        result = response.json()
+        logger.debug(f"Create conversation result: {result}")
+ 
+        if 'result' not in result or 'deploymentConversationId' not in result['result']:
+            l#ogger.error(f"Unexpected response structure: {result}")
+            raise HTTPException(status_code=401, detail="Invalid Abacus apikey")
+ 
+        return result["result"]["deploymentConversationId"], deployment_details["deploymentId"]
+    except Exception as e:
+        logger.error(f"Error creating conversation: {str(e)}", exc_info=True)
+        raise HTTPException(status_code=500, detail=f"Error creating conversation: {str(e)}")
+ 
+def serialize_openai_messages(messages):
+    def get_content(message):
+        try:
+            # Check if the 'content' key exists in message
+            if 'content' not in message:
+                return ''
+            if not isinstance(message['content'], list):
+                return message['content']
+            return message['content'][0]['text']
+        except KeyError as e:
+            raise HTTPException(status_code=400, detail="Invalid request body")
+ 
+    serialized_messages = [
+        f"{msg['role'].capitalize()}: {get_content(msg)}"
+        for msg in messages
+    ]
+ 
+    result = "\n\n".join(serialized_messages)
+ 
+    result += "Assistant: {...}\n\n"
+ 
+    return result.strip()
+ 
+CHAT_OUTPUT_PREFIX = 'data: {"id":"0","object":"0","created":0,"model":"0","choices":[{"index":0,"delta":{"content":'
+CHAT_OUTPUT_SUFFIX = '}}]}\n\n'
+ENDING_CHUNK = 'data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1694268190,"model":"gpt-4","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n'
+ 
+NS_PREFIX = '{"id":"chatcmpl-123","object":"chat.completion","created":1694268190,"model":"gpt-4","choices":[{"index":0,"message":{"role":"assistant","content":"'
+NS_SUFFIX = '"},"logprobs":null,"finish_reason":"stop"}],"usage":{"prompt_tokens":0,"completion_tokens":0,"total_tokens":0},"system_fingerprint":"0"}\n\n'
+ 
+async def stream_chat(apikey: str, conversation_id: str, body: Any, sse_flag=True):   
+    model = body["model"]
+    messages = body["messages"]
+ 
+    request_id = str(uuid.uuid4())
+    ws_url = f"wss://{BASE_HOST}/api/ws/chatLLMSendMessage?requestId={request_id}&docInfos=%5B%5D&deploymentConversationId={conversation_id}&llmName={model}&orgHost=apps"
+ 
+    headers = {
+        "apiKey": apikey,
+        "Origin": f"https://{BASE_HOST}",
+    }
+ 
+    if sse_flag:
+        data_prefix, data_suffix = CHAT_OUTPUT_PREFIX, CHAT_OUTPUT_SUFFIX
+        _Jd = jsonDumps
+    else:
+        data_prefix, data_suffix = "", ""
+        _Jd = lambda x: jsonDumps(x)[1:-1]
+        yield NS_PREFIX
+ 
+    try:
+        async with websockets.connect(ws_url, extra_headers=headers) as websocket:
+            serialized_msgs = serialize_openai_messages(messages)
+            await websocket.send(jsonDumps({"message": serialized_msgs}))
+            logger.debug(f"Sent message to WebSocket: {serialized_msgs}")
+ 
+            async for response in websocket:
+                logger.debug(f"Received WebSocket response: {response}")
+                data = json.loads(response)
+ 
+                if "segment" in data:
+                    segment = data['segment']
+                    if data['type'] == "image_url":
+                        segment = f"\n![Image]({segment})"
+                    yield data_prefix
+                    yield _Jd(segment)
+                    yield data_suffix
+                elif data.get("end", False):
+                    break
+ 
+        yield (ENDING_CHUNK if sse_flag else NS_SUFFIX)
+    except Exception as e:
+        logger.error(f"Error in WebSocket communication: {str(e)}", exc_info=True)
+        raise HTTPException(status_code=500, detail=f"WebSocket error: {str(e)}")
+ 
+async def handle_chat_completion(request: Request):
+    try:
+        body = await request.json()
+        logger.debug(f"Received request body: {body}")
+ 
+        auth_header = request.headers.get("Authorization")
+ 
+        if not auth_header or not auth_header.startswith("Bearer "):
+            raise HTTPException(status_code=401, detail="Invalid Authorization header")
+ 
+        abacus_token = auth_header[7:]  # Remove "Bearer " prefix
+ 
+        if not abacus_token:
+            raise HTTPException(status_code=401, detail="Empty Authorization token")
+ 
+        apikey = random.choice(abacus_token.split("|") or [abacus_token]) \
+                  if ("|" in abacus_token) \
+                  else abacus_token
+ 
+        apikey = convert_unicode_escape(apikey.strip())
+        logger.debug(f"Parsed apikey: {apikey}")
+ 
+        conversation_id, deployment_id = await create_conversation(apikey)
+        logger.debug(f"Created conversation with ID: {conversation_id}")
+ 
+        sse_flag = body.get("stream", (True if not "3.5" in body["model"] else False))
+ 
+        llm_name = map_model(body.get("model", ""))
+        body["model"] = llm_name
+        logger.info(f'#{deployment_id} - Querying {llm_name} in {("stream" if sse_flag else "non-stream")} mode')
+ 
+        return StreamingResponse(stream_chat(apikey, conversation_id, body, sse_flag), 
+                media_type=("text/event-stream" if sse_flag else "application/json") + \
+                    ";charset=UTF-8")
+    except Exception as e:
+        logger.error(f"Error in chat_completions: {str(e)}", exc_info=True)
+        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
+ 
+@app.post("/hf/v1/chat/completions")
+@limiter.limit(RATE_LIMIT)
+async def chat_completions(request: Request) -> StreamingResponse:
+    return await handle_chat_completion(request)
+ 
+def print_startup_info():
+    console = Console()
+    table = Table(title="Environment Variables & Available Models")
+ 
+    # Set up columns
+    table.add_column("Category", style="green")
+    table.add_column("Key", style="cyan")
+    table.add_column("Value", style="magenta")
+ 
+    # Add environment variables to the table
+    table.add_row("[bold]Environment Variables[/bold]", "", "")
+    for key, value in environment_variables.items():
+        table.add_row("", key, str(value))
+ 
+    # Add a separator row between the sections
+    table.add_row("", "", "")
+ 
+    # Add model mapping to the table
+    table.add_row("[bold]Available Models[/bold]", "", "")
+    for short_name, full_name in model_mapping.items():
+        table.add_row("", short_name, full_name)
+ 
+    # Print the table to the console
+    console.print(table)
+ 
+if __name__ == "__main__":
+    try:
+        import uvloop
+    except ImportError:
+        uvloop = None
+    if uvloop:
+        uvloop.install()
+ 
+    print_startup_info()
+ 
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=PORT, access_log=False)