add openai support

falkordb_gemini_kg/kg.py

@@ -5,6 +5,7 @@
 from falkordb_gemini_kg.classes.model_config import KnowledgeGraphModelConfig
 from falkordb_gemini_kg.steps.extract_data_step import ExtractDataStep
 from falkordb_gemini_kg.steps.graph_query_step import GraphQueryGenerationStep
+from falkordb_gemini_kg.fixtures.prompts import GRAPH_QA_SYSTEM, CYPHER_GEN_SYSTEM
 from falkordb_gemini_kg.steps.qa_step import QAStep
 from falkordb_gemini_kg.classes.ChatSession import ChatSession
 
@@ -124,7 +125,11 @@ def ask(self, question: str) -> str:
             >>> ans = kg.ask("List a few movies in which that actored played in", history)
         """
 
-        cypher_chat_session = self._model_config.cypher_generation.start_chat()
+        cypher_chat_session = (
+            self._model_config.cypher_generation.with_system_instruction(
+                CYPHER_GEN_SYSTEM.replace("#ONTOLOGY", str(self.ontology.to_json())),
+            ).start_chat()
+        )
         cypher_step = GraphQueryGenerationStep(
             ontology=self.ontology,
             chat_session=cypher_chat_session,
@@ -133,7 +138,9 @@ def ask(self, question: str) -> str:
 
         (context, cypher) = cypher_step.run(question)
 
-        qa_chat_session = self._model_config.qa.start_chat()
+        qa_chat_session = self._model_config.qa.with_system_instruction(
+            GRAPH_QA_SYSTEM
+        ).start_chat()
         qa_step = QAStep(
             chat_session=qa_chat_session,
         )

falkordb_gemini_kg/models/openai.py

@@ -0,0 +1,109 @@
+from .model import *
+from openai import OpenAI, completions
+
+
+class OpenAiGenerativeModel(GenerativeModel):
+
+    client: OpenAI = None
+
+    def __init__(
+        self,
+        model_name: str,
+        generation_config: GenerativeModelConfig | None = None,
+        system_instruction: str | None = None,
+    ):
+        self.model_name = model_name
+        self.generation_config = generation_config
+        self.system_instruction = system_instruction
+
+    def _get_model(self) -> OpenAI:
+        if self.client is None:
+            self.client = OpenAI()
+
+        return self.client
+
+    def with_system_instruction(self, system_instruction: str) -> "GenerativeModel":
+        self.system_instruction = system_instruction
+        self.client = None
+        self._get_model()
+
+        return self
+
+    def start_chat(self, args: dict | None = None) -> GenerativeModelChatSession:
+        return OpenAiChatSession(self, args)
+
+    def ask(self, message: str) -> GenerationResponse:
+        response = self.client.chat.completions.create(
+            model=self.model_name,
+            messages=[
+                {"role": "system", "content": self.system_instruction},
+                {"role": "user", "content": message},
+            ],
+            max_tokens=self.generation_config.max_output_tokens,
+            temperature=self.generation_config.temperature,
+            top_p=self.generation_config.top_p,
+            top_k=self.generation_config.top_k,
+            stop=self.generation_config.stop_sequences,
+        )
+        return self.parse_generate_content_response(response)
+
+    def parse_generate_content_response(self, response: any) -> GenerationResponse:
+        return GenerationResponse(
+            text=response.choices[0].message.content,
+            finish_reason=(
+                FinishReason.STOP
+                if response.choices[0].finish_reason == "stop"
+                else (
+                    FinishReason.MAX_TOKENS
+                    if response.choices[0].finish_reason == "length"
+                    else FinishReason.OTHER
+                )
+            ),
+        )
+
+
+class OpenAiChatSession(GenerativeModelChatSession):
+
+    _history = []
+
+    def __init__(self, model: OpenAiGenerativeModel, args: dict | None = None):
+        self._model = model
+        self._args = args
+        self._history = (
+            [{"role": "system", "content": self._model.system_instruction}]
+            if self._model.system_instruction is not None
+            else []
+        )
+
+    def send_message(self, message: str) -> GenerationResponse:
+        prompt = []
+        prompt.extend(self._history)
+        prompt.append({"role": "user", "content": message})
+        response = self._model.client.chat.completions.create(
+            model=self._model.model_name,
+            messages=prompt,
+            max_tokens=(
+                self._model.generation_config.max_output_tokens
+                if self._model.generation_config is not None
+                else None
+            ),
+            temperature=(
+                self._model.generation_config.temperature
+                if self._model.generation_config is not None
+                else None
+            ),
+            top_p=(
+                self._model.generation_config.top_p
+                if self._model.generation_config is not None
+                else None
+            ),
+            stop=(
+                self._model.generation_config.stop_sequences
+                if self._model.generation_config is not None
+                else None
+            ),
+        )
+        content = self._model.parse_generate_content_response(response)
+        self._history.append({"role": "user", "content": message})
+        self._history.append({"role": "system", "content": content.text})
+        return content

pyproject.toml

@@ -21,6 +21,7 @@ backoff = "^2.2.1"
 python-abc = "^0.2.0"
 ratelimit = "^2.2.1"
 python-dotenv = "^1.0.1"
+openai = "^1.35.9"
 
 [tool.poetry.group.test.dependencies]
 pytest = "^8.2.1"

tests/test_kg.py → tests/test_kg_gemini.py

@@ -20,7 +20,7 @@
 vertexai.init(project=os.getenv("PROJECT_ID"), location=os.getenv("REGION"))
 
 
-class TestKG(unittest.TestCase):
+class TestKGGemini(unittest.TestCase):
     """
     Test Knowledge Graph
     """