Update pytests

.github/workflows/test.yml

@@ -24,6 +24,10 @@ jobs:
           - 11434:11434
         volumes:
           - ollama:/root/.ollama
+
+    strategy:
+      matrix:
+        model: [gemini/gemini-1.5-flash-001, openai/gpt-4o] # List of models
 
     steps:
       - uses: actions/checkout@v4
@@ -73,4 +77,6 @@ jobs:
           REGION: ${{ vars.REGION }}
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
           GOOGLE_API_KEY: ${{ secrets.GOOGLE_API_KEY }}
+          GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
+          TEST_MODEL: ${{ matrix.model }} # Pass the model as an environment variable
         run: poetry run pytest

graphrag_sdk/fixtures/prompts.py

@@ -433,12 +433,9 @@
 
 
 For example, given the question "Which managers own Neo4j stocks?", the OpenCypher statement should look like this:
-```
 MATCH (m:Manager)-[:OWNS]->(s:Stock)
 WHERE s.name CONTAINS 'Neo4j'
-RETURN m, s
-```
-"""
+RETURN m, s"""
 
 CYPHER_GEN_PROMPT = """
 Using the ontology provided, generate an OpenCypher statement to query the graph database returning all relevant entities, relationships, and attributes to answer the question below.
@@ -447,7 +444,7 @@
 Please think if your answer is a valid Cypher query, and correct it if it is not.
 
 Question: {question}
-"""
+Your generated Cypher: """
 
 
 CYPHER_GEN_PROMPT_WITH_ERROR = """
@@ -472,9 +469,8 @@
 Respect the order of the relationships; the arrows should always point from the "source" to the "target".
 
 Last Answer: {last_answer}
-
 Question: {question}
-"""
+Your generated Cypher: """
 
 GRAPH_QA_SYSTEM = """
 You are an assistant that helps to form nice and human understandable answers.

graphrag_sdk/models/litellm.py

@@ -74,7 +74,7 @@ def check_valid_key(self, model: str):
                 model=model, messages=messages, max_tokens=10
             )
             return True
-        except:
+        except Exception as e:
             return False
 
     def check_and_pull_model(self) -> None: