add-litellm

Author: galshubeli

graphrag_sdk/models/litellm.py

@@ -0,0 +1,134 @@
+from typing import Optional
+from litellm import completion
+from .model import (
+    OutputMethod,
+    GenerativeModel,
+    GenerativeModelConfig,
+    GenerationResponse,
+    FinishReason,
+    GenerativeModelChatSession,
+)
+
+
+class LiteLLMGenerativeModel(GenerativeModel):
+    """
+    A generative model that interfaces with the LiteLLM for chat completions.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        generation_config: Optional[GenerativeModelConfig] = None,
+        system_instruction: Optional[str] = None,
+    ):
+        self.model_name = model_name
+        self.generation_config = generation_config or GenerativeModelConfig()
+        self.system_instruction = system_instruction
+
+
+    def with_system_instruction(self, system_instruction: str) -> "GenerativeModel":
+        self.system_instruction = system_instruction
+        return self
+
+    def start_chat(self, args: Optional[dict] = None) -> GenerativeModelChatSession:
+        return LiteLLMChatSession(self, args)
+
+    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
+                )
+            ),
+        )
+
+    def to_json(self) -> dict:
+        return {
+            "model_name": self.model_name,
+            "generation_config": self.generation_config.to_json(),
+            "system_instruction": self.system_instruction,
+        }
+
+    @staticmethod
+    def from_json(json: dict) -> "GenerativeModel":
+        return LiteLLMGenerativeModel(
+            json["model_name"],
+            generation_config=GenerativeModelConfig.from_json(
+                json["generation_config"]
+            ),
+            system_instruction=json["system_instruction"],
+        )
+
+
+class LiteLLMChatSession(GenerativeModelChatSession):
+
+    _history = []
+
+    def __init__(self, model: LiteLLMGenerativeModel, args: Optional[dict] = 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, output_method: OutputMethod = OutputMethod.DEFAULT) -> GenerationResponse:
+        generation_config = self._get_generation_config(output_method)
+        prompt = []
+        prompt.extend(self._history)
+        prompt.append({"role": "user", "content": message[:14385]})
+        response = completion(
+            model=self._model.model_name,
+            messages=prompt,
+            **generation_config
+            )
+        content = self._model.parse_generate_content_response(response)
+        self._history.append({"role": "user", "content": message})
+        self._history.append({"role": "assistant", "content": content.text})
+        return content
+    
+    def _get_generation_config(self, output_method: OutputMethod):
+        config = self._model.generation_config.to_json()
+        if output_method == OutputMethod.JSON:
+            config['temperature'] = 0
+            config['response_format'] = { "type": "json_object" }
+        
+        return config
+    
+    def delete_last_message(self):
+        """
+        Deletes the last message exchange (user message and assistant response) from the chat history.
+        Preserves the system message if present.
+        
+        Example:
+            Before:
+            [
+                {"role": "system", "content": "System message"},
+                {"role": "user", "content": "User message"},
+                {"role": "assistant", "content": "Assistant response"},
+            ]
+            After:
+            [
+                {"role": "system", "content": "System message"},
+            ]
+
+        Note: Does nothing if the chat history is empty or contains only a system message.
+        """
+        # Keep at least the system message if present
+        min_length = 1 if self._model.system_instruction else 0
+        if len(self._history) - 2 >= min_length:
+            self._history.pop()
+            self._history.pop()
+        else:
+            # Reset to initial state with just system message if present
+            self._history = (
+            [{"role": "system", "content": self._model.system_instruction}]
+            if self._model.system_instruction is not None
+            else []
+        )

tests/test_kg_litellm_openai.py

@@ -0,0 +1,102 @@
+import re
+import logging
+import unittest
+from falkordb import FalkorDB
+from dotenv import load_dotenv
+from graphrag_sdk.entity import Entity
+from graphrag_sdk.source import Source
+from graphrag_sdk.relation import Relation
+from graphrag_sdk.ontology import Ontology
+from graphrag_sdk.attribute import Attribute, AttributeType
+from graphrag_sdk.models.litellm import LiteLLMGenerativeModel
+from graphrag_sdk import KnowledgeGraph, KnowledgeGraphModelConfig
+
+load_dotenv()
+
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
+class TestKGLiteLLM(unittest.TestCase):
+    """
+    Test Knowledge Graph
+    """
+
+    @classmethod
+    def setUpClass(cls):
+
+        cls.ontology = Ontology([], [])
+
+        cls.ontology.add_entity(
+            Entity(
+                label="Actor",
+                attributes=[
+                    Attribute(
+                        name="name",
+                        attr_type=AttributeType.STRING,
+                        unique=True,
+                        required=True,
+                    ),
+                ],
+            )
+        )
+        cls.ontology.add_entity(
+            Entity(
+                label="Movie",
+                attributes=[
+                    Attribute(
+                        name="title",
+                        attr_type=AttributeType.STRING,
+                        unique=True,
+                        required=True,
+                    ),
+                ],
+            )
+        )
+        cls.ontology.add_relation(
+            Relation(
+                label="ACTED_IN",
+                source="Actor",
+                target="Movie",
+                attributes=[
+                    Attribute(
+                        name="role",
+                        attr_type=AttributeType.STRING,
+                        unique=False,
+                        required=False,
+                    ),
+                ],
+            )
+        )
+        cls.graph_name = "IMDB_openai"
+        model = LiteLLMGenerativeModel(model_name="gpt-4o")
+        cls.kg = KnowledgeGraph(
+            name=cls.graph_name,
+            ontology=cls.ontology,
+            model_config=KnowledgeGraphModelConfig.with_model(model),
+        )
+
+    def test_kg_creation(self):
+
+        file_path = "tests/data/madoff.txt"
+
+        sources = [Source(file_path)]
+
+        self.kg.process_sources(sources)
+
+        chat = self.kg.chat_session()
+        answer = chat.send_message("How many actors acted in a movie?")
+        answer = answer['response']
+
+        logger.info(f"Answer: {answer}")
+
+        actors_count = re.findall(r'\d+', answer)
+        num_actors = 0 if len(actors_count) == 0 else int(actors_count[0])
+
+        assert num_actors > 10, "The number of actors found should be greater than 10"
+
+    def test_kg_delete(self):
+        self.kg.delete()
+
+        db = FalkorDB()
+        graphs = db.list_graphs()
+        self.assertNotIn(self.graph_name, graphs)