chore(update): Update the License Year + Fixes mypy errors (#18)

* chore: mypy initial fixes

* chore: poetry config change + spanner import type ignore

* chore: isort fix

* chore(update): Update the License Year

* chore(update):updates the workers for integratrion test

* chore(update):updates the workers for integratrion test

* chore(update): update to auto workers

* chore(test):fix test failures

* chore(update): update to auto workers

* chore(update): update coverage for workers to run parallely

* chore(update):fixing the workers issue

* chore(test):update command to test

* chore(revert):revert the changes done for auto workers

* chore(fix): fix the MALFORMED_FUNCTION_CALL in the tests results

* chore(test):update the test case

* chore(fix): fix my-py errors

* chore(fix): fix the lint

* chore(fix):fix the MALFORMED_FUNCTION_CALL by Gemini

* chore(fix):fix old changes replicas

* chore(fix):fix lint errors

---------

Co-authored-by: gRedHeadphone <[email protected]>

Author: shreya6goel

.coveragerc

@@ -1,4 +1,5 @@
 [run]
+parallel = true
 branch = true
 omit =
     */__init__.py

CHANGELOG.md

@@ -1,3 +1 @@
-# Changelog
-
-## 0.1.0 (2025-07-01)
+# Changelog

integration.cloudbuild.yaml

@@ -26,7 +26,7 @@ steps:
   - id: Run integration tests
     name: 'python:${_VERSION}'
     entrypoint: python
-    args: ["-m", "pytest", "-n", "3", "--cov=llama_index_spanner", "--cov-config=.coveragerc", "tests/"]
+    args: ["-m", "pytest", "-n", "5", "--cov=llama_index_spanner", "--cov-config=.coveragerc", "tests/"]
     secretEnv: ["GOOGLE_API_KEY"]
     env:
       - "PROJECT_ID=$PROJECT_ID"

pyproject.toml

@@ -10,7 +10,8 @@ authors = [
 ]
 dependencies = [
     "google-cloud-spanner>=3.41.0, <4.0.0",
-    "llama-index-core==0.12.41"
+    "llama-index-core>=0.12.48, <1.0.0",
+    "llama-index>=0.12.41, <1.0.0"
 ]
 classifiers = [
     "Intended Audience :: Developers",
@@ -46,8 +47,8 @@ test = [
     "mypy==1.16.0",
     "pytest==8.3.3",
     "pytest-xdist==3.6.1",
-    "pytest-asyncio==0.24.0",
     "pytest-cov==5.0.0",
+    "pytest-asyncio==0.24.0",
     "llama-index-readers-wikipedia==0.3.0",
     "llama-index-embeddings-google-genai==0.1.0",
     "llama-index-llms-google-genai==0.1.4",

requirements.txt

@@ -14,4 +14,4 @@
 
 google-cloud-spanner==3.51.0
 llama-index-core==0.12.48
-llama-index==0.12.41
+llama-index==0.12.48

src/llama_index_spanner/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Google LLC
+# Copyright 2025 Google LLC
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

src/llama_index_spanner/graph_retriever.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Google LLC
+# Copyright 2025 Google LLC
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -23,7 +23,7 @@
 from llama_index.core.prompts import PromptType
 from llama_index.core.retrievers import CustomPGRetriever, VectorContextRetriever
 from llama_index.core.schema import NodeWithScore, QueryBundle, TextNode
-from llama_index.core.vector_stores.types import VectorStore
+from llama_index.core.vector_stores.types import BasePydanticVectorStore
 from pydantic import BaseModel
 
 from .graph_utils import extract_gql, fix_gql_syntax
@@ -78,7 +78,6 @@ def __init__(
         graph_store: SpannerPropertyGraphStore,
         llm: Optional[LLM] = None,
         text_to_gql_prompt: Optional[PromptTemplate] = None,
-        response_template: Optional[str] = None,
         gql_validator: Optional[Callable[[str], bool]] = None,
         include_raw_response_as_metadata: Optional[bool] = False,
         max_gql_fix_retries: Optional[int] = 1,
@@ -93,7 +92,6 @@ def __init__(
           graph_store: The SpannerPropertyGraphStore to query.
           llm: The LLM to use.
           text_to_gql_prompt: The prompt to use for generating the GQL query.
-          response_template: The template to use for formatting the response.
           gql_validator: A function to validate the GQL query.
           include_raw_response_as_metadata: If true, includes the raw response as
             metadata.
@@ -179,7 +177,7 @@ def calculate_score_for_predicted_response(
         gql_response_score = self.llm.predict(
             GQL_RESPONSE_SCORING_TEMPLATE, question=question, retrieved_context=response
         )
-        return gql_response_score
+        return float(gql_response_score.strip())
 
     def retrieve_from_graph(
         self, query_bundle: schema.QueryBundle
@@ -208,16 +206,19 @@ def retrieve_from_graph(
 
         # 2. Verify gql query using LLM
         if self.verify_gql:
-            verify_response = self.llm.predict(
-                GQL_VERIFY_PROMPT,
-                question=question,
-                generated_gql=generated_gql,
-                schema=schema_str,
-                format_instructions=GQL_VERIFY_PROMPT.output_parser.format_string,
-            )
+            if GQL_VERIFY_PROMPT.output_parser:
+                verify_response = self.llm.predict(
+                    GQL_VERIFY_PROMPT,
+                    question=question,
+                    generated_gql=generated_gql,
+                    schema=schema_str,
+                    format_instructions=GQL_VERIFY_PROMPT.output_parser.format,
+                )
 
-            output_parser = verify_gql_output_parser.parse(verify_response)
-            verified_gql = fix_gql_syntax(output_parser.verified_gql)
+                output_parser = verify_gql_output_parser.parse(verify_response)
+                verified_gql = fix_gql_syntax(output_parser.verified_gql)
+            else:
+                raise ValueError("GQL_VERIFY_PROMPT is missing its output_parser.")
         else:
             verified_gql = generated_gql
 
@@ -259,7 +260,7 @@ def retrieve_from_graph(
     async def aretrieve_from_graph(
         self, query_bundle: QueryBundle
     ) -> List[NodeWithScore]:
-        return await self.retrieve_from_graph(query_bundle)
+        return self.retrieve_from_graph(query_bundle)
 
 
 class SpannerGraphCustomRetriever(CustomPGRetriever):
@@ -269,13 +270,12 @@ def init(
         self,
         ## vector context retriever params
         embed_model: Optional[BaseEmbedding] = None,
-        vector_store: Optional[VectorStore] = None,
+        vector_store: Optional[BasePydanticVectorStore] = None,
         similarity_top_k: int = 4,
         path_depth: int = 2,
         ## text-to-gql params
         llm_text_to_gql: Optional[LLM] = None,
         text_to_gql_prompt: Optional[PromptTemplate] = None,
-        response_template: Optional[str] = None,
         gql_validator: Optional[Callable[[str], bool]] = None,
         include_raw_response_as_metadata: Optional[bool] = False,
         max_gql_fix_retries: Optional[int] = 1,
@@ -297,7 +297,6 @@ def init(
           path_depth: The depth of the path to retrieve.
           llm_text_to_gql: The LLM to use for text to GQL conversion.
           text_to_gql_prompt: The prompt to use for generating the GQL query.
-          response_template: The template to use for formatting the response.
           gql_validator: A function to validate the GQL query.
           include_raw_response_as_metadata: Whether to include the raw response as
             metadata.
@@ -311,6 +310,12 @@ def init(
           llmranker_top_n: The number of top nodes to return.
           **kwargs: Additional keyword arguments.
         """
+
+        if not isinstance(self._graph_store, SpannerPropertyGraphStore):
+            raise TypeError(
+                "SpannerGraphCustomRetriever requires a SpannerPropertyGraphStore."
+            )
+
         self.llm = llm_text_to_gql or Settings.llm
         if self.llm is None:
             raise ValueError("`llm for Text to GQL` cannot be none")
@@ -328,7 +333,6 @@ def init(
             graph_store=self._graph_store,
             llm=llm_text_to_gql,
             text_to_gql_prompt=text_to_gql_prompt,
-            response_template=response_template,
             gql_validator=gql_validator,
             include_raw_response_as_metadata=include_raw_response_as_metadata,
             max_gql_fix_retries=max_gql_fix_retries,
@@ -342,7 +346,7 @@ def init(
             top_n=llmranker_top_n,
         )
 
-    def generate_synthesized_response(self, question: str, response: str) -> float:
+    def generate_synthesized_response(self, question: str, response: str) -> str:
         gql_synthesized_response = self.llm.predict(
             GQL_SYNTHESIS_RESPONSE_TEMPLATE,
             question=question,

src/llama_index_spanner/graph_utils.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Google LLC
+# Copyright 2025 Google LLC
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

src/llama_index_spanner/prompts.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Google LLC
+# Copyright 2025 Google LLC
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

src/llama_index_spanner/property_graph_store.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Google LLC
+# Copyright 2025 Google LLC
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.