fix-excep-process-fail

Author: galshubeli

graphrag_sdk/models/litellm.py

@@ -25,6 +25,7 @@ def __init__(
         model_name: str,
         generation_config: Optional[GenerativeModelConfig] = None,
         system_instruction: Optional[str] = None,
+        host: Optional[str] = None,
     ):
         """
         Initialize the LiteModel with the required parameters.
@@ -33,9 +34,13 @@ def __init__(
             model_name (str): The name and the provider for the LiteLLM client.
             generation_config (Optional[GenerativeModelConfig]): Configuration settings for generation.
             system_instruction (Optional[str]): Instruction to guide the model.
+            host (Optional[str]): Host for connecting to the Ollama API.
         """
 
+        self.host = host
         self.model = model_name
+        
+        # LiteLLM model name format: <provider>/<model_name> - Example: openai/gpt-4o
         if "/" in model_name:
             self.provider = model_name.split("/")[0]
             self.model_name = model_name.split("/")[1]
@@ -67,7 +72,7 @@ def credentials_validation(self) -> None:
                 raise ValueError("Missing Gemini API key in the environment.")
 
         elif self.provider == "ollama":
-            self.ollama_client = Client()
+            self.ollama_client = Client(self.host)
             self.check_and_pull_model()
 
     def check_and_pull_model(self) -> None:
@@ -224,11 +229,11 @@ def send_message(self, message: str, output_method: OutputMethod = OutputMethod.
             response = completion(
                 model=self._model.model,
                 messages=self._chat_history,
+                api_base=self._model.host,
                 **generation_config
             )
         except Exception as e:
-            # Handle exception (e.g., log error, retry, or raise a custom exception)
-            raise RuntimeError("Error during completion request") from e
+            raise ValueError(f"Error during completion request, please check the credentials - {e}")
         content = self._model.parse_generate_content_response(response)
         self._chat_history.append({"role": "assistant", "content": content.text})
         return content

graphrag_sdk/steps/create_ontology_step.py

@@ -17,7 +17,6 @@
 from graphrag_sdk.models import (
     GenerativeModel,
     GenerativeModelChatSession,
-    GenerativeModelConfig,
     GenerationResponse,
     FinishReason,
 )
@@ -71,7 +70,12 @@ def run(self, boundaries: Optional[str] = None):
                 tasks.append(task)
 
             # Wait for all tasks to complete
-            wait(tasks)
+            done, _ = wait(tasks)  # Get completed tasks
+            for task in done:
+                try:
+                    task.result()  # Re-raise any exceptions from _process_source
+                except Exception as e:
+                    raise ValueError(f"Error in task: {e}")
 
         for task in tasks:
             self.ontology = self.ontology.merge_with(task.result())

graphrag_sdk/steps/extract_data_step.py

@@ -90,7 +90,12 @@ def run(self, instructions: str = None):
                 tasks.append(task)
 
             # Wait for all tasks to complete
-            wait(tasks)
+            done, _ = wait(tasks)  # Get completed tasks
+            for task in done:
+                try:
+                    task.result()  # Re-raise any exceptions from _process_source
+                except Exception as e:
+                    raise ValueError(f"Error in task: {e}")
 
     def _process_source(
         self,
@@ -195,7 +200,7 @@ def _process_source(
 
         except Exception as e:
             logger.exception(e)
-            raise e
+            raise ValueError(f"Error processing source: {e}")
 
     def _create_entity(self, graph: Graph, args: dict, ontology: Ontology):
         # Get unique attributes from entity

tests/test_kg_litellm_openai.py

@@ -84,19 +84,12 @@ def test_kg_creation(self):
         self.kg.process_sources(sources)
 
         chat = self.kg.chat_session()
-        answer = chat.send_message("How many actors acted in a movie?")
+        answer = chat.send_message("How many different actors there are?")
         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)
+        assert num_actors > 10, "The number of actors found should be greater than 10"