wip deepseek r1

sia/llm_engine/openai_llm_engine.py

@@ -0,0 +1,75 @@
+from typing import Callable, Iterator
+import openai
+import tiktoken
+
+from . import LlmEngine
+
+class OpenAILlmEngine(LlmEngine):
+    """
+    LLM Engine implementation using OpenAI's API.
+    Supports streaming responses from chat completion models.
+    """
+    
+    def __init__(
+        self,
+        model: str,
+        temperature: float,
+        token_limit: int,
+        api_key: str,
+    ):
+        """
+        Initialize the OpenAI LLM Engine.
+
+        Args:
+            model: OpenAI model to use
+            temperature: Temperature for sampling
+            api_key: OpenAI API key
+            token_limit: Maximum number of tokens to generate
+        """
+        self._model = model
+        self._temperature = temperature
+        self._token_limit = token_limit
+
+        self._client = openai.Client(
+            api_key=api_key,
+        )
+                        
+    def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context}
+        ]
+        
+        stream = self._client.chat.completions.create(
+            model=self._model,
+            messages=messages,
+            temperature=self._temperature,
+            stream=True,
+        )
+
+        try:
+            for chunk in stream:
+                if should_stop():
+                    break
+                if content := chunk.choices[0].delta.content:
+                    yield content
+        finally:
+            stream.close()
+            #stream.response.close()
+    
+    def token_count(self, system_prompt: str, main_context: str) -> int:
+        """
+        Calculate the total token count for the system prompt and context.
+        
+        Args:
+            system_prompt: The system prompt string
+            main_context: The main context string
+            
+        Returns:
+            int: Total number of tokens
+        """
+        encoding = tiktoken.encoding_for_model(self._model)
+        return len(encoding.encode(system_prompt)) + len(encoding.encode(main_context))
+    
+    def token_limit(self) -> int:
+        return self._token_limit