Added openai llm engine

requirements.txt

@@ -1,6 +1,7 @@
 accelerate
 aiohttp
 bs4
+openai
 python-dotenv
 torch
 transformers

sia/__main__.py

@@ -5,11 +5,11 @@ import asyncio
 import mimetypes
 import time
 
-
 from .config import Config
 from .hf_llm_engine import HfLlmEngine
 from .llm_engine import LlmEngine
 from .local_llm_engine import LocalLlmEngine
+from .openai_llm_engine import OpenAILlmEngine
 from .response_parser import ResponseParser
 from .system_metrics import SystemMetrics
 from .web_agent import WebAgent
@@ -50,6 +50,11 @@ class Main:
                     model_id=self._config.model,
                     api_token=self._config.api_token
                 )
+            case "openai":
+                self._llm = OpenAILlmEngine(
+                    model=self._config.model,
+                    api_key=self._config.api_token
+                )
             case "test":
                 self._llm = TestLLM()
             case _:

sia/openai_llm_engine.py

@@ -0,0 +1,57 @@
+from typing import Iterator, Optional
+import openai
+import json
+
+from .llm_engine import LlmEngine
+
+class OpenAILlmEngine(LlmEngine):
+    """
+    LLM Engine implementation using OpenAI's API.
+    Supports streaming responses from chat completion models.
+    """
+    
+    def __init__(
+        self,
+        model: str,
+        api_key: str,
+    ):
+        """
+        Initialize the OpenAI LLM Engine.
+
+        Args:
+            model: OpenAI model to use (default: gpt-4)
+            api_key: OpenAI API key. If None, will try to read from OPENAI_API_KEY env var
+        """
+        self._model = model
+        
+        # Initialize OpenAI client
+        self.client = openai.Client(
+            api_key=api_key,
+        )
+                        
+    def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
+        """
+        Run inference using the system prompt and main context.
+
+        Args:
+            system_prompt: The system prompt string
+            main_context: The main context string after templating
+
+        Returns:
+            Iterator[str]: An iterator that yields the generated text in chunks.
+        """
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context}
+        ]
+        
+        stream = self.client.chat.completions.create(
+            model=self._model,
+            messages=messages,
+            stream=True,
+            temperature=0.3,
+        )
+
+        for chunk in stream:
+            if content := chunk.choices[0].delta.content:
+                yield content