Create deepinfra_handler.py

deepinfra_handler.py

@@ -0,0 +1,65 @@
+import json
+import requests
+from typing import Dict, Any, Generator, Optional
+
+class DeepInfraHandler:
+    API_URL = "https://api.deepinfra.com/v1/openai/chat/completions"
+    
+    def __init__(self):
+        self.headers = {
+            "Accept": "text/event-stream",
+            "Accept-Encoding": "gzip, deflate, br, zstd",
+            "Content-Type": "application/json",
+            "Connection": "keep-alive",
+            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36",
+        }
+
+    def _prepare_payload(self, **kwargs) -> Dict[str, Any]:
+        """Prepare the payload for the API request"""
+        return {
+            "model": kwargs.get("model"),
+            "messages": kwargs.get("messages"),
+            "temperature": kwargs.get("temperature", 0.7),
+            "max_tokens": kwargs.get("max_tokens", 4096),
+            "top_p": kwargs.get("top_p", 1.0),
+            "frequency_penalty": kwargs.get("frequency_penalty", 0.0),
+            "presence_penalty": kwargs.get("presence_penalty", 0.0),
+            "stop": kwargs.get("stop", []),
+            "stream": kwargs.get("stream", False)
+        }
+
+    def generate_completion(self, **kwargs) -> Any:
+        """Generate completion based on streaming preference"""
+        payload = self._prepare_payload(**kwargs)
+        
+        response = requests.post(
+            self.API_URL,
+            headers=self.headers,
+            json=payload,
+            stream=payload["stream"]
+        )
+        
+        if payload["stream"]:
+            return self._handle_streaming_response(response)
+        return self._handle_regular_response(response)
+
+    def _handle_streaming_response(self, response) -> Generator[str, None, None]:
+        """Handle streaming response from the API"""
+        for line in response.iter_lines(decode_unicode=True):
+            if line.startswith("data:"):
+                try:
+                    content = json.loads(line[5:])
+                    if content == "[DONE]":
+                        continue
+                    delta_content = content.get("choices", [{}])[0].get("delta", {}).get("content")
+                    if delta_content:
+                        yield delta_content
+                except:
+                    continue
+
+    def _handle_regular_response(self, response) -> Dict[str, Any]:
+        """Handle regular (non-streaming) response from the API"""
+        try:
+            return response.json()
+        except Exception as e:
+            raise Exception(f"Error processing response: {str(e)}")