Create v1.py

v1.py

@@ -0,0 +1,114 @@
+import requests
+import json
+from typing import Generator
+import os
+from dotenv import load_dotenv
+import re
+
+# Load environment variables from .env file
+load_dotenv()
+
+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)