Create main.py

main.py

@@ -0,0 +1,428 @@
+import os
+import json
+import time
+import uuid
+from typing import List, Dict, Optional, Union, Generator, Any
+
+# --- Core Dependencies ---
+import uvicorn
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+from pydantic import BaseModel, Field
+from curl_cffi.requests import Session
+from curl_cffi import CurlError
+
+# --- Environment Configuration ---
+QODO_API_KEY = os.getenv("QODO_API_KEY", "useme")
+QODO_URL = os.getenv("QODO_URL", "https://hello.com")
+QODO_INFO_URL = os.getenv("QODO_INFO_URL", "https://openai.com")
+
+# --- Recreated/Mocked webscout Dependencies ---
+# This section recreates the necessary classes and functions
+# to make the QodoAI provider self-contained.
+
+# webscout.exceptions
+class exceptions:
+    class FailedToGenerateResponseError(Exception):
+        pass
+
+# webscout.AIutel.sanitize_stream
+def sanitize_stream(data: Generator[bytes, None, None], content_extractor: callable, **kwargs: Any) -> Generator[str, None, None]:
+    """
+    Parses a stream of byte chunks, extracts complete JSON objects,
+    and yields content processed by the content_extractor.
+    """
+    buffer = ""
+    for byte_chunk in data:
+        buffer += byte_chunk.decode('utf-8', errors='ignore')
+        
+        start_index = 0
+        while True:
+            # Find the start of a potential JSON object
+            try:
+                obj_start = buffer.index('{', start_index)
+            except ValueError:
+                # No more objects in buffer, keep the remainder for the next chunk
+                buffer = buffer[start_index:]
+                break
+
+            # Find the corresponding end brace
+            brace_count = 1
+            i = obj_start + 1
+            while i < len(buffer) and brace_count > 0:
+                if buffer[i] == '{':
+                    brace_count += 1
+                elif buffer[i] == '}':
+                    brace_count -= 1
+                i += 1
+            
+            if brace_count == 0:  # Found a complete object
+                json_str = buffer[obj_start:i]
+                try:
+                    json_obj = json.loads(json_str)
+                    content = content_extractor(json_obj)
+                    if content:
+                        yield content
+                except json.JSONDecodeError:
+                    pass  # Skip malformed JSON
+                start_index = i  # Move past the processed object
+            else:
+                # Incomplete object, wait for more data
+                buffer = buffer[start_index:]
+                break
+
+# webscout.Provider.OPENAI.utils (Pydantic Models)
+class Tool(BaseModel):
+    type: str = "function"
+    function: Dict[str, Any]
+
+class ChatCompletionMessage(BaseModel):
+    role: str
+    content: Optional[str] = None
+    tool_calls: Optional[List[Dict]] = None
+
+class Choice(BaseModel):
+    index: int
+    message: Optional[ChatCompletionMessage] = None
+    finish_reason: Optional[str] = None
+    delta: Optional[Dict] = Field(default_factory=dict)
+
+class ChoiceDelta(BaseModel):
+    content: Optional[str] = None
+    role: Optional[str] = None
+
+class ChoiceStreaming(BaseModel):
+    index: int
+    delta: ChoiceDelta
+    finish_reason: Optional[str] = None
+
+class CompletionUsage(BaseModel):
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+class ChatCompletion(BaseModel):
+    id: str
+    choices: List[Choice]
+    created: int
+    model: str
+    object: str = "chat.completion"
+    usage: CompletionUsage
+
+class ChatCompletionChunk(BaseModel):
+    id: str
+    choices: List[ChoiceStreaming]
+    created: int
+    model: str
+    object: str = "chat.completion.chunk"
+    usage: Optional[CompletionUsage] = None
+
+# webscout.Provider.OPENAI.base
+class BaseCompletions:
+    def __init__(self, client: Any):
+        self._client = client
+
+class BaseChat:
+    def __init__(self, client: Any):
+        self.completions = Completions(client)
+
+class OpenAICompatibleProvider:
+    def __init__(self, **kwargs: Any):
+        pass
+
+# Attempt to import LitAgent, fallback if not available
+try:
+    from webscout.litagent import LitAgent
+except ImportError:
+    LitAgent = None
+
+# --- QodoAI Provider Code (from the prompt) ---
+
+class Completions(BaseCompletions):
+    def create(
+        self,
+        *,
+        model: str,
+        messages: List[Dict[str, Any]],
+        stream: bool = False,
+        **kwargs: Any
+    ) -> Union[ChatCompletion, Generator[ChatCompletionChunk, None, None]]:
+        """
+        Creates a model response for the given chat conversation.
+        Mimics openai.chat.completions.create
+        """
+        user_prompt = ""
+        for message in reversed(messages):
+            if message.get("role") == "user":
+                user_prompt = message.get("content", "")
+                break
+
+        if not user_prompt:
+            raise ValueError("No user message found in messages")
+
+        payload = self._client._build_payload(user_prompt, model)
+        payload["stream"] = stream
+        payload["custom_model"] = model
+
+        request_id = f"chatcmpl-{uuid.uuid4()}"
+        created_time = int(time.time())
+
+        if stream:
+            return self._create_stream(request_id, created_time, model, payload, user_prompt)
+        else:
+            return self._create_non_stream(request_id, created_time, model, payload, user_prompt)
+
+    def _create_stream(
+        self, request_id: str, created_time: int, model: str, payload: Dict[str, Any], user_prompt: str
+    ) -> Generator[ChatCompletionChunk, None, None]:
+        try:
+            response = self._client.session.post(
+                self._client.url,
+                json=payload,
+                stream=True,
+                timeout=self._client.timeout,
+                impersonate="chrome110"
+            )
+
+            if response.status_code == 401:
+                raise exceptions.FailedToGenerateResponseError("Invalid Qodo API key provided.")
+            elif response.status_code != 200:
+                raise IOError(f"Qodo request failed with status code {response.status_code}: {response.text}")
+
+            prompt_tokens = len(user_prompt.split())
+            completion_tokens = 0
+            
+            processed_stream = sanitize_stream(
+                data=response.iter_content(chunk_size=None),
+                content_extractor=QodoAI._qodo_extractor
+            )
+
+            for content_chunk in processed_stream:
+                if content_chunk:
+                    completion_tokens += len(content_chunk.split())
+                    
+                    delta = ChoiceDelta(content=content_chunk, role="assistant")
+                    choice = ChoiceStreaming(index=0, delta=delta, finish_reason=None)
+                    chunk = ChatCompletionChunk(id=request_id, choices=[choice], created=created_time, model=model)
+                    yield chunk
+
+            final_choice = ChoiceStreaming(index=0, delta=ChoiceDelta(), finish_reason="stop")
+            yield ChatCompletionChunk(id=request_id, choices=[final_choice], created=created_time, model=model)
+
+        except CurlError as e:
+            raise exceptions.FailedToGenerateResponseError(f"Request failed (CurlError): {e}")
+        except Exception as e:
+            raise exceptions.FailedToGenerateResponseError(f"An unexpected error occurred ({type(e).__name__}): {e}")
+
+    def _create_non_stream(
+        self, request_id: str, created_time: int, model: str, payload: Dict[str, Any], user_prompt: str
+    ) -> ChatCompletion:
+        try:
+            payload["stream"] = False
+            response = self._client.session.post(
+                self._client.url,
+                json=payload,
+                timeout=self._client.timeout,
+                impersonate="chrome110"
+            )
+            
+            if response.status_code == 401:
+                raise exceptions.FailedToGenerateResponseError("Invalid Qodo API key provided.")
+            elif response.status_code != 200:
+                raise IOError(f"Qodo request failed with status code {response.status_code}: {response.text}")
+
+            response_text = response.text
+            full_response = ""
+            
+            # This logic parses concatenated JSON objects from the response body.
+            current_json = ""
+            brace_count = 0
+            json_objects = []
+            lines = response_text.strip().split('\n')
+            for line in lines:
+                current_json += line
+                brace_count += line.count('{') - line.count('}')
+                if brace_count == 0 and current_json:
+                    json_objects.append(current_json)
+                    current_json = ""
+            
+            for json_str in json_objects:
+                try:
+                    json_obj = json.loads(json_str)
+                    content = QodoAI._qodo_extractor(json_obj)
+                    if content:
+                        full_response += content
+                except json.JSONDecodeError:
+                    pass
+
+            prompt_tokens = len(user_prompt.split())
+            completion_tokens = len(full_response.split())
+            total_tokens = prompt_tokens + completion_tokens
+
+            message = ChatCompletionMessage(role="assistant", content=full_response)
+            choice = Choice(index=0, message=message, finish_reason="stop")
+            usage = CompletionUsage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens)
+            return ChatCompletion(id=request_id, choices=[choice], created=created_time, model=model, usage=usage)
+
+        except CurlError as e:
+            raise exceptions.FailedToGenerateResponseError(f"Request failed (CurlError): {e}")
+        except Exception as e:
+            raise exceptions.FailedToGenerateResponseError(f"Request failed ({type(e).__name__}): {e}")
+
+class Chat(BaseChat):
+    def __init__(self, client: 'QodoAI'):
+        self.completions = Completions(client)
+
+class QodoAI(OpenAICompatibleProvider):
+    AVAILABLE_MODELS = ["gpt-4.1", "gpt-4o", "o3", "o4-mini", "claude-4-sonnet", "gemini-2.5-pro"]
+
+    def __init__(self, api_key: str, **kwargs: Any):
+        super().__init__(api_key=api_key, **kwargs)
+        
+        self.url = QODO_URL
+        self.info_url = QODO_INFO_URL
+        self.timeout = 600
+        self.api_key = api_key
+        
+        self.user_agent = "axios/1.10.0"
+        self.session_id = self._get_session_id()
+        self.request_id = str(uuid.uuid4())
+        
+        self.headers = {
+            "Accept": "text/plain", "Accept-Encoding": "gzip, deflate, br, zstd",
+            "Accept-Language": "en-US,en;q=0.9", "Authorization": f"Bearer {self.api_key}",
+            "Connection": "close", "Content-Type": "application/json",
+            "host": "api.cli.qodo.ai", "Request-id": self.request_id,
+            "Session-id": self.session_id, "User-Agent": self.user_agent,
+        }
+        
+        self.session = Session()
+        self.session.headers.update(self.headers)
+        self.chat = Chat(self)
+        
+    @staticmethod
+    def _qodo_extractor(chunk: Union[str, Dict[str, Any]]) -> Optional[str]:
+        if isinstance(chunk, dict):
+            data = chunk.get("data", {})
+            if isinstance(data, dict):
+                tool_args = data.get("tool_args", {})
+                if isinstance(tool_args, dict) and "content" in tool_args:
+                    return tool_args.get("content")
+                if "content" in data:
+                    return data["content"]
+        return None
+
+    def _get_session_id(self) -> str:
+        try:
+            temp_session = Session()
+            temp_headers = {
+                "Authorization": f"Bearer {self.api_key}",
+                "User-Agent": self.user_agent,
+            }
+            temp_session.headers.update(temp_headers)
+            
+            response = temp_session.get(self.info_url, timeout=self.timeout, impersonate="chrome110")
+            
+            if response.status_code == 200:
+                return response.json().get("session-id", f"fallback-{uuid.uuid4()}")
+            elif response.status_code == 401:
+                raise exceptions.FailedToGenerateResponseError("Invalid Qodo API key. Please check your QODO_API_KEY environment variable.")
+            else:
+                raise exceptions.FailedToGenerateResponseError(f"Failed to get session_id from Qodo: HTTP {response.status_code}")
+        except Exception as e:
+            raise exceptions.FailedToGenerateResponseError(f"Failed to connect to Qodo API to get session_id: {e}")
+
+    def _build_payload(self, prompt: str, model: str) -> Dict[str, Any]:
+        return {
+            "agent_type": "cli", "session_id": self.session_id,
+            "user_data": {"extension_version": "0.7.2", "os_platform": "win32"},
+            "tools": {"web_search": []}, "user_request": prompt,
+            "execution_strategy": "act", "custom_model": model, "stream": True
+        }
+        
+# --- FastAPI Application ---
+
+app = FastAPI(
+    title="QodoAI OpenAI-Compatible API",
+    description="Provides an OpenAI-compatible interface for the QodoAI service.",
+    version="1.0.0"
+)
+
+# Initialize the client at startup
+try:
+    client = QodoAI(api_key=QODO_API_KEY)
+except exceptions.FailedToGenerateResponseError as e:
+    print(f"FATAL: Could not initialize QodoAI client: {e}")
+    print("Please ensure the QODO_API_KEY environment variable is set correctly.")
+    client = None
+
+# --- API Models ---
+
+class Model(BaseModel):
+    id: str
+    object: str = "model"
+    created: int = Field(default_factory=lambda: int(time.time()))
+    owned_by: str = "qodoai"
+
+class ModelList(BaseModel):
+    object: str = "list"
+    data: List[Model]
+
+class ChatCompletionRequest(BaseModel):
+    model: str
+    messages: List[Dict[str, Any]]
+    max_tokens: Optional[int] = 2049
+    stream: bool = False
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+    tools: Optional[List[Dict[str, Any]]] = None
+    tool_choice: Optional[str] = None
+    
+# --- API Endpoints ---
+
+@app.on_event("startup")
+async def startup_event():
+    if client is None:
+        # This will prevent the app from starting if the client failed to init
+        raise RuntimeError("QodoAI client could not be initialized. Check API key and connectivity.")
+    print("QodoAI client initialized successfully.")
+
+@app.get("/v1/models", response_model=ModelList)
+async def list_models():
+    """Lists the available models from the QodoAI provider."""
+    model_data = [Model(id=model_id) for model_id in QodoAI.AVAILABLE_MODELS]
+    return ModelList(data=model_data)
+
+@app.post("/v1/chat/completions")
+async def create_chat_completion(request: ChatCompletionRequest):
+    """Creates a chat completion, supporting both streaming and non-streaming modes."""
+    if client is None:
+        raise HTTPException(status_code=500, detail="QodoAI client is not available.")
+        
+    params = request.model_dump(exclude_none=True)
+    
+    try:
+        if request.stream:
+            async def stream_generator():
+                try:
+                    generator = client.chat.completions.create(**params)
+                    for chunk in generator:
+                        yield f"data: {chunk.model_dump_json()}\n\n"
+                    yield "data: [DONE]\n\n"
+                except exceptions.FailedToGenerateResponseError as e:
+                    error_payload = {"error": {"message": str(e), "type": "api_error"}}
+                    yield f"data: {json.dumps(error_payload)}\n\n"
+                    yield "data: [DONE]\n\n"
+
+            return StreamingResponse(stream_generator(), media_type="text/event-stream")
+        else:
+            response = client.chat.completions.create(**params)
+            return JSONResponse(content=response.model_dump())
+
+    except exceptions.FailedToGenerateResponseError as e:
+        raise HTTPException(status_code=500, detail=str(e))
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)