Create app.py

app.py

@@ -0,0 +1,113 @@
+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
+
+load_dotenv()
+
+app = FastAPI()
+
+class v1:
+    """
+    A class to interact with the Prefind AI API.
+    """
+
+    AVAILABLE_MODELS = ["llama", "claude"]
+
+    def __init__(
+        self,
+        model: str = "claude",
+        timeout: int = 30,
+        proxies: dict = {},
+    ):
+        """
+        Initializes the Prefind 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}"
+            )
+
+        for line in response.iter_lines(decode_unicode=True):
+            if line:
+                if line.startswith("data: "):
+                    data_str = line[6:]
+                    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:
+                                yield content
+
+    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"):
+    if model not in v1.AVAILABLE_MODELS:
+        raise HTTPException(status_code=400, detail=f"Model '{model}' is not supported. Choose from {v1.AVAILABLE_MODELS}.")
+
+    ai = v1(model=model)
+    
+    def response_generator():
+        for chunk in ai.chat(prompt):
+            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=7860)