From 229f4f4ac7994b5a12abfcde9c702e75a8aacb2b Mon Sep 17 00:00:00 2001
From: geens <niels.geens@gmail.com>
Date: Wed, 13 Nov 2024 11:46:29 +0100
Subject: [PATCH] Added openai llm engine

---
 requirements.txt         |  1 +
 sia/__main__.py          |  7 ++++-
 sia/openai_llm_engine.py | 57 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 64 insertions(+), 1 deletion(-)
 create mode 100644 sia/openai_llm_engine.py

diff --git a/requirements.txt b/requirements.txt
index 07d56b5..da16518 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,7 @@
 accelerate
 aiohttp
 bs4
+openai
 python-dotenv
 torch
 transformers
\ No newline at end of file
diff --git a/sia/__main__.py b/sia/__main__.py
index 8338bcc..eace4e6 100644
--- a/sia/__main__.py
+++ b/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 _:
diff --git a/sia/openai_llm_engine.py b/sia/openai_llm_engine.py
new file mode 100644
index 0000000..88607a8
--- /dev/null
+++ b/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
\ No newline at end of file