Update main.py

main.py

@@ -2,6 +2,7 @@ import os
 import json
 import time
 import uuid
+import sys
 from typing import List, Dict, Optional, Union, Generator, Any
 
 # --- Core Dependencies ---
@@ -13,84 +14,98 @@ 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")
+QODO_API_KEY = os.getenv("QODO_API_KEY") # No default key to encourage setting it explicitly
+QODO_URL = os.getenv("QODO_URL", "https://*")
+QODO_INFO_URL = os.getenv("QODO_INFO_URL", "*")
 
-# --- Recreated/Mocked webscout Dependencies ---
+# --- Recreated/Mocked webscout & OpenAI 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')
+        obj_start_indices = [i for i, char in enumerate(buffer) if char == '{']
+        if not obj_start_indices:
+            continue
         
         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:
+        for obj_start in obj_start_indices:
+            if obj_start < start_index:
+                continue
+            
+            brace_count = 0
+            obj_end = -1
+            for i in range(obj_start, len(buffer)):
                 if buffer[i] == '{':
                     brace_count += 1
                 elif buffer[i] == '}':
                     brace_count -= 1
-                i += 1
+                if brace_count == 0:
+                    obj_end = i
+                    break
             
-            if brace_count == 0:  # Found a complete object
-                json_str = buffer[obj_start:i]
+            if obj_end != -1:
+                json_str = buffer[obj_start:obj_end + 1]
                 try:
                     json_obj = json.loads(json_str)
                     content = content_extractor(json_obj)
                     if content:
                         yield content
+                    start_index = obj_end + 1
                 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
+                    continue
+        buffer = buffer[start_index:]
+
+# --- OpenAI-Compatible Pydantic Models ---
+
+# Request Models
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+    name: Optional[str] = None
+    tool_calls: Optional[List[Dict]] = None
+    tool_call_id: Optional[str] = None
+
+class Function(BaseModel):
+    name: str
+    description: Optional[str] = None
+    parameters: Dict[str, Any]
 
-# webscout.Provider.OPENAI.utils (Pydantic Models)
 class Tool(BaseModel):
     type: str = "function"
-    function: Dict[str, Any]
+    function: Function
 
+class ChatCompletionRequest(BaseModel):
+    model: str
+    messages: List[ChatMessage]
+    max_tokens: Optional[int] = 2049
+    stream: bool = False
+    temperature: Optional[float] = 1.0
+    top_p: Optional[float] = 1.0
+    tools: Optional[List[Tool]] = None
+    tool_choice: Optional[Union[str, Dict]] = None
+
+# Response Models
 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 Choice(BaseModel):
+    index: int
+    message: ChatCompletionMessage
+    finish_reason: Optional[str] = "stop"
+
 class ChoiceStreaming(BaseModel):
     index: int
     delta: ChoiceDelta
@@ -117,7 +132,8 @@ class ChatCompletionChunk(BaseModel):
     object: str = "chat.completion.chunk"
     usage: Optional[CompletionUsage] = None
 
-# webscout.Provider.OPENAI.base
+
+# --- Base Provider Structure ---
 class BaseCompletions:
     def __init__(self, client: Any):
         self._client = client
@@ -130,13 +146,7 @@ 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) ---
+# --- QodoAI Provider Logic ---
 
 class Completions(BaseCompletions):
     def create(
@@ -147,22 +157,23 @@ class Completions(BaseCompletions):
         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
-        """
+        
+        # Warn about unsupported parameters
+        unsupported_params = ['temperature', 'top_p', 'tools', 'tool_choice', 'max_tokens']
+        for param in unsupported_params:
+            if param in kwargs:
+                print(f"Warning: Parameter '{param}' is not supported by the QodoAI provider and will be ignored.", file=sys.stderr)
+
         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")
+            raise ValueError("No user message with 'role': 'user' 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())
@@ -172,102 +183,45 @@ class Completions(BaseCompletions):
         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]:
+    def _create_stream(self, request_id, created_time, model, payload, user_prompt) -> 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}")
+            with self._client.session.post(self._client.url, json=payload, stream=True, timeout=self._client.timeout, impersonate="chrome110") as response:
+                if response.status_code == 401:
+                    raise exceptions.FailedToGenerateResponseError("Invalid Qodo API key provided.")
+                response.raise_for_status()
+
+                for content_chunk in sanitize_stream(response.iter_content(chunk_size=8192), QodoAI._qodo_extractor):
+                    if content_chunk:
+                        delta = ChoiceDelta(content=content_chunk, role="assistant")
+                        choice = ChoiceStreaming(index=0, delta=delta, finish_reason=None)
+                        yield ChatCompletionChunk(id=request_id, choices=[choice], created=created_time, model=model)
+                
+                final_delta = ChoiceDelta()
+                final_choice = ChoiceStreaming(index=0, delta=final_delta, finish_reason="stop")
+                yield ChatCompletionChunk(id=request_id, choices=[final_choice], created=created_time, model=model)
         except Exception as e:
-            raise exceptions.FailedToGenerateResponseError(f"An unexpected error occurred ({type(e).__name__}): {e}")
+            raise exceptions.FailedToGenerateResponseError(f"Stream generation failed: {e}")
 
-    def _create_non_stream(
-        self, request_id: str, created_time: int, model: str, payload: Dict[str, Any], user_prompt: str
-    ) -> ChatCompletion:
+    def _create_non_stream(self, request_id, created_time, model, payload, user_prompt) -> ChatCompletion:
         try:
             payload["stream"] = False
-            response = self._client.session.post(
-                self._client.url,
-                json=payload,
-                timeout=self._client.timeout,
-                impersonate="chrome110"
-            )
-            
+            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.raise_for_status()
 
-            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
+            full_response = "".join(list(sanitize_stream(iter([response.content]), QodoAI._qodo_extractor)))
 
             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)
+            usage = CompletionUsage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_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}")
+            raise exceptions.FailedToGenerateResponseError(f"Non-stream generation failed: {e}")
 
 class Chat(BaseChat):
     def __init__(self, client: 'QodoAI'):
@@ -278,26 +232,14 @@ class QodoAI(OpenAICompatibleProvider):
 
     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.url, self.info_url, self.timeout, self.api_key = QODO_URL, QODO_INFO_URL, 600, 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,
+            "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json",
+            "User-Agent": self.user_agent, "Session-id": self.session_id
         }
-        
-        self.session = Session()
-        self.session.headers.update(self.headers)
+        self.session = Session(headers=self.headers)
         self.chat = Chat(self)
         
     @staticmethod
@@ -305,24 +247,13 @@ class QodoAI(OpenAICompatibleProvider):
         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"]
+                content = data.get("content") or (data.get("tool_args", {}) or {}).get("content")
+                if content: return 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")
-            
+            response = Session(headers={"Authorization": f"Bearer {self.api_key}"}).get(self.info_url, timeout=self.timeout)
             if response.status_code == 200:
                 return response.json().get("session-id", f"fallback-{uuid.uuid4()}")
             elif response.status_code == 401:
@@ -333,13 +264,8 @@ class QodoAI(OpenAICompatibleProvider):
             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
-        }
-        
+        return {"agent_type": "cli", "session_id": self.session_id, "user_request": prompt, "custom_model": model, "stream": True}
+
 # --- FastAPI Application ---
 
 app = FastAPI(
@@ -348,57 +274,34 @@ app = FastAPI(
     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 ---
+client: Optional[QodoAI] = None
 
 @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.")
+def startup_event():
+    global client
+    if not QODO_API_KEY:
+        raise RuntimeError("QODO_API_KEY environment variable not set. The server cannot start without an API key.")
+    try:
+        client = QodoAI(api_key=QODO_API_KEY)
+        print("QodoAI client initialized successfully.")
+    except exceptions.FailedToGenerateResponseError as e:
+        raise RuntimeError(f"FATAL: Could not initialize QodoAI client: {e}")
 
-@app.get("/v1/models", response_model=ModelList)
+@app.get("/v1/models", response_model_exclude_none=True)
 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)
+    models_data = [
+        {"id": model_id, "object": "model", "created": int(time.time()), "owned_by": "qodoai"}
+        for model_id in QodoAI.AVAILABLE_MODELS
+    ]
+    return {"object": "list", "data": models_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.")
-        
+        raise HTTPException(status_code=503, detail="QodoAI client is not available or failed to initialize.")
+    
     params = request.model_dump(exclude_none=True)
     
     try:
@@ -408,12 +311,11 @@ async def create_chat_completion(request: ChatCompletionRequest):
                     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"}}
+                    error_payload = {"error": {"message": str(e), "type": "api_error", "code": 500}}
                     yield f"data: {json.dumps(error_payload)}\n\n"
+                finally:
                     yield "data: [DONE]\n\n"
-
             return StreamingResponse(stream_generator(), media_type="text/event-stream")
         else:
             response = client.chat.completions.create(**params)
@@ -425,4 +327,7 @@ async def create_chat_completion(request: ChatCompletionRequest):
         raise HTTPException(status_code=400, detail=str(e))
 
 if __name__ == "__main__":
+    if not QODO_API_KEY:
+        print("Error: The QODO_API_KEY environment variable must be set.", file=sys.stderr)
+        sys.exit(1)
     uvicorn.run(app, host="0.0.0.0", port=8000)