Update app.py

app.py

@@ -1,126 +1,13 @@
-import requests
-import json
-from typing import Generator
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
 import uvicorn
-from dotenv import load_dotenv
-import os
-import re
-
-# Load environment variables from .env file
-load_dotenv()
+from v1 import v1
+from v2 import v2
 
 app = FastAPI()
 
-class v1:
-    """ 
-    A class to interact with the v1 AI API.
-    """
-
-    AVAILABLE_MODELS = ["llama", "claude"]
-
-    def __init__(
-        self,
-        model: str = "claude",
-        timeout: int = 300,
-        proxies: dict = {},
-    ):
-        """
-        Initializes the v1 AI API with given parameters.
-        Args:
-            model (str, optional): The AI model to use for text generation. Defaults to "claude". 
-                                    Options: "llama", "claude".
-            timeout (int, optional): Http request timeout. Defaults to 30.
-            proxies (dict, optional): Http request proxies. Defaults to {}.
-        """
-        if model not in self.AVAILABLE_MODELS:
-            raise ValueError(f"Model '{model}' is not supported. Choose from {self.AVAILABLE_MODELS}.")
-
-        self.session = requests.Session()
-        self.api_endpoint = os.getenv("API_ENDPOINT")
-        self.timeout = timeout
-        self.model = model
-        self.device_token = self.get_device_token()
-
-        self.session.headers.update(
-            {
-                "Content-Type": "application/json",
-                "Accept": "text/event-stream",
-            }
-        )
-        self.session.proxies = proxies
-
-    def get_device_token(self) -> str:
-        device_token_url = os.getenv("DEVICE_TOKEN_URL")
-        headers = {"Content-Type": "application/json; charset=utf-8"}
-        data = {}
-        response = requests.post(
-            device_token_url, headers=headers, data=json.dumps(data)
-        )
-
-        if response.status_code == 200:
-            device_token_data = response.json()
-            return device_token_data["sessionToken"]
-        else:
-            raise Exception(
-                f"Failed to get device token - ({response.status_code}, {response.reason}) - {response.text}"
-            )
-
-    def ask(self, prompt: str) -> Generator[str, None, None]:
-        search_data = {"query": prompt, "deviceToken": self.device_token}
-
-        response = self.session.post(
-            self.api_endpoint, json=search_data, stream=True, timeout=self.timeout
-        )
-        if not response.ok:
-            raise Exception(
-                f"Failed to generate response - ({response.status_code}, {response.reason}) - {response.text}"
-            )
-
-        buffer = ""
-        for line in response.iter_lines(decode_unicode=True):
-            if line:
-                if line.startswith("data: "):
-                    data_str = line[6:]
-                    try:
-                        data = json.loads(data_str)
-                        if data['type'] == 'chunk':
-                            model = data['model']
-                            if (self.model == "llama" and model == 'OPENROUTER_LLAMA_3') or \
-                               (self.model == "claude" and model == 'OPENROUTER_CLAUDE'):
-                                content = data['chunk']['content']
-                                if content:
-                                    buffer += content
-                                    # Check if we have a complete line or paragraph
-                                    lines = buffer.split('\n')
-                                    if len(lines) > 1:
-                                        for complete_line in lines[:-1]:
-                                            yield self.format_text(complete_line) + '\n'
-                                        buffer = lines[-1]
-                    except KeyError:
-                        pass
-                    except json.JSONDecodeError:
-                        pass
-
-        # Yield any remaining content in the buffer
-        if buffer:
-            yield self.format_text(buffer)
-
-        yield "[DONE]"
-
-    def format_text(self, text: str) -> str:
-        # Convert *text* to <i>text</i> for italic
-        text = re.sub(r'\*(.*?)\*', r'<i>\1</i>', text)
-        return text
-
-    def chat(self, prompt: str) -> Generator[str, None, None]:
-        """Stream responses as string chunks"""
-        return self.ask(prompt)
-
-
 @app.get("/Search/pro")
-async def chat(prompt: str, model: str = "claude"):
+async def v1_chat(prompt: str, model: str = "claude"):
     if model not in v1.AVAILABLE_MODELS:
         raise HTTPException(status_code=400, detail=f"Model '{model}' is not supported. Choose from {v1.AVAILABLE_MODELS}.")
 
@@ -132,6 +19,15 @@ async def chat(prompt: str, model: str = "claude"):
     
     return StreamingResponse(response_generator(), media_type="text/event-stream")
 
+@app.get("/v2/search")
+async def v2_chat(prompt: str):
+    ai = v2()
+    
+    def response_generator():
+        for chunk in ai.chat(prompt, stream=True):
+            yield f"data: {chunk}\n\n"
+    
+    return StreamingResponse(response_generator(), media_type="text/event-stream")
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=8000)