2 introduce llm model abstraction

.github/workflows/spellcheck.yml

@@ -0,0 +1,17 @@
+name: spellcheck
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+jobs:
+  check-spelling:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+      - name: Check Spelling
+        uses: rojopolis/spellcheck-github-actions@0.33.1
+        with:
+          config_path: .spellcheck.yml
+          task_name: Markdown

.github/workflows/test.yml

@@ -0,0 +1,58 @@
+name: Run tests
+
+on:
+  push:
+  pull_request:
+
+jobs:
+  test:
+    permissions:
+      contents: "read"
+      id-token: "write"
+    runs-on: ubuntu-latest
+    services:
+      falkordb:
+        image: falkordb/falkordb:latest
+        ports:
+          - 6379:6379
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: 3.x # Update with desired Python version
+
+      - name: Cache Poetry virtualenv
+        id: cache
+        uses: actions/cache@v4
+        with:
+          path: ~/.poetry/virtualenvs
+          key: ${{ runner.os }}-poetry-${{ hashFiles('poetry.lock') }}
+          restore-keys: |
+            ${{ runner.os }}-poetry-
+
+      - id: "auth"
+        uses: "google-github-actions/auth@v2"
+        with:
+          credentials_json: ${{ secrets.GCP_SA_KEY }}
+
+      - name: "Set up Cloud SDK"
+        uses: "google-github-actions/setup-gcloud@v2"
+        with:
+          version: ">= 363.0.0"
+
+      - name: Install Poetry
+        if: steps.cache.outputs.cache-hit != true
+        run: |
+          curl -sSL https://install.python-poetry.org | python3 -
+
+      - name: Install dependencies
+        run: poetry install
+
+      - name: Run tests
+        env:
+          PROJECT_ID: ${{ secrets.PROJECT_ID }}
+          REGION: ${{ vars.REGION }}
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+        run: poetry run pytest

falkordb_gemini_kg/__init__.py

@@ -1,7 +1,7 @@
 from .classes.source import Source
 from .classes.ontology import Ontology
 from .kg import KnowledgeGraph
-from .classes.model_config import KnowledgeGraphModelConfig, StepModelConfig
+from .classes.model_config import KnowledgeGraphModelConfig
 from .steps.create_ontology_step import CreateOntologyStep
 
 # Setup Null handler

falkordb_gemini_kg/classes/__init__.py

@@ -2,8 +2,10 @@
 from .source import Source
 from .node import Node
 from .edge import Edge
+from .attribute import Attribute, AttributeType
 
 # Setup Null handler
 import logging
+
 logger = logging.getLogger(__name__)
 logger.addHandler(logging.NullHandler())

falkordb_gemini_kg/classes/attribute.py

@@ -5,23 +5,23 @@
 logger = logging.getLogger(__name__)
 
 
-class _AttributeType:
+class AttributeType:
     STRING = "string"
     NUMBER = "number"
     BOOLEAN = "boolean"
 
     @staticmethod
     def fromString(txt: str):
         if txt.isdigit():
-            return _AttributeType.NUMBER
+            return AttributeType.NUMBER
         elif txt.lower() in ["true", "false"]:
-            return _AttributeType.BOOLEAN
-        return _AttributeType.STRING
+            return AttributeType.BOOLEAN
+        return AttributeType.STRING
 
 
 class Attribute:
     def __init__(
-        self, name: str, attr_type: _AttributeType, unique: bool, required: bool = False
+        self, name: str, attr_type: AttributeType, unique: bool, required: bool = False
     ):
         self.name = name
         self.type = attr_type
@@ -32,9 +32,9 @@ def __init__(
     def from_json(txt: str):
         txt = txt if isinstance(txt, dict) else json.loads(txt)
         if txt["type"] not in [
-            _AttributeType.STRING,
-            _AttributeType.NUMBER,
-            _AttributeType.BOOLEAN,
+            AttributeType.STRING,
+            AttributeType.NUMBER,
+            AttributeType.BOOLEAN,
         ]:
             raise Exception(f"Invalid attribute type: {txt['type']}")
         return Attribute(
@@ -52,9 +52,9 @@ def from_string(txt: str):
         required = "*" in txt
 
         if attr_type not in [
-            _AttributeType.STRING,
-            _AttributeType.NUMBER,
-            _AttributeType.BOOLEAN,
+            AttributeType.STRING,
+            AttributeType.NUMBER,
+            AttributeType.BOOLEAN,
         ]:
             raise Exception(f"Invalid attribute type: {attr_type}")
 

falkordb_gemini_kg/classes/edge.py

@@ -1,7 +1,7 @@
 import json
 import re
 import logging
-from .attribute import Attribute, _AttributeType
+from .attribute import Attribute, AttributeType
 from falkordb import Node as GraphNode, Edge as GraphEdge
 from falkordb_gemini_kg.fixtures.regex import *
 
@@ -58,7 +58,7 @@ def from_graph(edge: GraphEdge, nodes: list[GraphNode]):
             [
                 Attribute(
                     attr,
-                    _AttributeType.fromString(edge.properties),
+                    AttributeType.fromString(edge.properties),
                     "!" in edge.properties[attr],
                     "*" in edge.properties[attr],
                 )

falkordb_gemini_kg/classes/model_config.py

@@ -1,66 +1,22 @@
-from vertexai.generative_models import GenerationConfig
-
-
-class StepModelGenerationConfig:
-    def __init__(
-        self,
-        temperature: float,
-        top_p: float,
-        top_k: int,
-        candidate_count: int,
-        max_output_tokens: int,
-        stop_sequences: list[str],
-    ):
-        self.temperature = temperature
-        self.top_p = top_p
-        self.top_k = top_k
-        self.candidate_count = candidate_count
-        self.max_output_tokens = max_output_tokens
-        self.stop_sequences = stop_sequences
-
-    def to_generation_config(self) -> GenerationConfig:
-        return GenerationConfig(
-            temperature=self.temperature,
-            top_p=self.top_p,
-            top_k=self.top_k,
-            candidate_count=self.candidate_count,
-            max_output_tokens=self.max_output_tokens,
-            stop_sequences=self.stop_sequences,
-        )
-
-
-class StepModelConfig:
-
-    def __init__(
-        self, model: str, generation_config: StepModelGenerationConfig | None = None
-    ):
-        self.model = model
-        self.generation_config = generation_config
+from falkordb_gemini_kg.models import GenerativeModel
 
 
 class KnowledgeGraphModelConfig:
 
     def __init__(
         self,
-        extract_data: StepModelConfig | None = None,
-        cypher_generation: StepModelConfig | None = None,
-        qa: StepModelConfig | None = None,
+        extract_data: GenerativeModel,
+        cypher_generation: GenerativeModel,
+        qa: GenerativeModel,
     ):
         self.extract_data = extract_data
         self.cypher_generation = cypher_generation
         self.qa = qa
 
     @staticmethod
-    def from_dict(d: dict):
-        model = d.get("model")
-        generation_config = d.get("generation_config")
-        extract_data = StepModelConfig(model=model, generation_config=generation_config)
-        cypher_generation = StepModelConfig(
-            model=model, generation_config=generation_config
-        )
-        qa = StepModelConfig(model=model, generation_config=generation_config)
+    def with_model(model: GenerativeModel):
         return KnowledgeGraphModelConfig(
-            extract_data=extract_data,
-            cypher_generation=cypher_generation,
-            qa=qa,
+            extract_data=model,
+            cypher_generation=model,
+            qa=model,
         )

falkordb_gemini_kg/classes/node.py

@@ -1,6 +1,6 @@
 import json
 import logging
-from .attribute import Attribute, _AttributeType
+from .attribute import Attribute, AttributeType
 from falkordb import Node as GraphNode
 
 logger = logging.getLogger(__name__)
@@ -18,7 +18,7 @@ def from_graph(node: GraphNode):
             [
                 Attribute(
                     attr,
-                    _AttributeType.fromString(node.properties[attr]),
+                    AttributeType.fromString(node.properties[attr]),
                     "!" in node.properties[attr],
                 )
                 for attr in node.properties

falkordb_gemini_kg/classes/ontology.py

@@ -1,7 +1,7 @@
 import json
 from falkordb import Graph
 from falkordb_gemini_kg.classes.source import AbstractSource
-from falkordb_gemini_kg.classes.model_config import StepModelConfig
+from falkordb_gemini_kg.models import GenerativeModel
 import falkordb_gemini_kg
 import logging
 from .edge import Edge
@@ -19,12 +19,12 @@ def __init__(self, nodes: list[Node] = [], edges: list[Edge] = []):
     def from_sources(
         sources: list[AbstractSource],
         boundaries: str,
-        model_config: StepModelConfig,
+        model: GenerativeModel,
     ) -> "Ontology":
         step = falkordb_gemini_kg.CreateOntologyStep(
             sources=sources,
             ontology=Ontology(),
-            model_config=model_config,
+            model=model,
         )
 
         return step.run(boundaries=boundaries)

falkordb_gemini_kg/helpers.py

@@ -165,9 +165,6 @@ def validate_cypher_edge_directions(cypher: str, ontology: Ontology):
             if ontology_edge is None:
                 errors.append(f"Edge {edge_label} not found in ontology")
 
-            logger.debug(
-                f"ontology_edge: {ontology_edge}"
-            )
             if (
                 not ontology_edge.source.label == source_label
                 or not ontology_edge.target.label == target_label

falkordb_gemini_kg/kg.py

@@ -104,7 +104,7 @@ def _create_graph_with_sources(self, sources: list[AbstractSource] | None = None
         step = ExtractDataStep(
             sources=list(sources),
             ontology=self.ontology,
-            model_config=self._model_config.extract_data,
+            model=self._model_config.extract_data,
             graph=self.graph,
         )
 
@@ -124,16 +124,18 @@ 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_step = GraphQueryGenerationStep(
             ontology=self.ontology,
-            model_config=self._model_config.cypher_generation,
+            chat_session=cypher_chat_session,
             graph=self.graph,
         )
 
         (context, cypher) = cypher_step.run(question)
 
+        qa_chat_session = self._model_config.qa.start_chat()
         qa_step = QAStep(
-            model_config=self._model_config.qa,
+            chat_session=qa_chat_session,
         )
 
         answer = qa_step.run(question, cypher, context)

falkordb_gemini_kg/models/__init__.py

@@ -0,0 +1 @@
+from .model import *