mypy

stellasia · stellasia · commit 13a3cc96cc69 · 2025-04-24T13:35:38.000+02:00
diff --git a/src/neo4j_graphrag/llm/vertexai_llm.py b/src/neo4j_graphrag/llm/vertexai_llm.py
@@ -251,14 +251,14 @@ def _to_tool_call(self, function_call: FunctionCall) -> ToolCall:
             arguments=function_call.args,
         )
 
-    def _parse_tool_response(self, response) -> ToolCallResponse:
+    def _parse_tool_response(self, response: GenerationResponse) -> ToolCallResponse:
         function_calls = response.candidates[0].function_calls
         return ToolCallResponse(
             tool_calls=[self._to_tool_call(f) for f in function_calls],
             content=None,
         )
 
-    def _parse_content_response(self, response) -> LLMResponse:
+    def _parse_content_response(self, response: GenerationResponse) -> LLMResponse:
         return LLMResponse(
             content=response.text,
         )

Original file line number	Diff line number	Diff line change
`@@ -251,14 +251,14 @@ def _to_tool_call(self, function_call: FunctionCall) -> ToolCall:`
`251`	`251`	`arguments=function_call.args,`
`252`	`252`	`)`
`253`	`253`
`254`		`- def _parse_tool_response(self, response) -> ToolCallResponse:`
	`254`	`+ def _parse_tool_response(self, response: GenerationResponse) -> ToolCallResponse:`
`255`	`255`	`function_calls = response.candidates[0].function_calls`
`256`	`256`	`return ToolCallResponse(`
`257`	`257`	`tool_calls=[self._to_tool_call(f) for f in function_calls],`
`258`	`258`	`content=None,`
`259`	`259`	`)`
`260`	`260`
`261`		`- def _parse_content_response(self, response) -> LLMResponse:`
	`261`	`+ def _parse_content_response(self, response: GenerationResponse) -> LLMResponse:`
`262`	`262`	`return LLMResponse(`
`263`	`263`	`content=response.text,`
`264`	`264`	`)`