kensho-kfinance 3.2.5__py3-none-any.whl → 3.2.7__py3-none-any.whl

kfinance/domains/prices/tests/test_price_tools.py

@@ -5,8 +5,10 @@ from kfinance.conftest import SPGI_TRADING_ITEM_ID
 from kfinance.domains.companies.company_models import COMPANY_ID_PREFIX
 from kfinance.domains.prices.price_tools import (
     GetHistoryMetadataFromIdentifiers,
+    GetHistoryMetadataFromIdentifiersResp,
     GetPricesFromIdentifiers,
     GetPricesFromIdentifiersArgs,
+    GetPricesFromIdentifiersResp,
 )
 from kfinance.integrations.tool_calling.tool_calling_models import ToolArgsWithIdentifiers
 
@@ -27,12 +29,14 @@ class TestGetHistoryMetadataFromIdentifiers:
             "instrument_type": "Equity",
             "symbol": "SPGI",
         }
-        expected_resp = {
-            "results": {"SPGI": metadata_resp},
-            "errors": [
-                "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
-            ],
-        }
+        expected_resp = GetHistoryMetadataFromIdentifiersResp.model_validate(
+            {
+                "results": {"SPGI": metadata_resp},
+                "errors": [
+                    "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
+                ],
+            }
+        )
 
         requests_mock.get(
             url=f"https://kfinance.kensho.com/api/v1/pricing/{SPGI_TRADING_ITEM_ID}/metadata",
@@ -80,33 +84,35 @@ class TestGetPricesFromIdentifiers:
             url=f"https://kfinance.kensho.com/api/v1/pricing/{SPGI_TRADING_ITEM_ID}/none/none/day/adjusted",
             json=self.prices_resp,
         )
-        expected_response = {
-            "results": {
-                "SPGI": {
-                    "prices": [
-                        {
-                            "date": "2024-04-11",
-                            "open": {"value": "424.26", "unit": "USD"},
-                            "high": {"value": "425.99", "unit": "USD"},
-                            "low": {"value": "422.04", "unit": "USD"},
-                            "close": {"value": "422.92", "unit": "USD"},
-                            "volume": {"value": "1129158", "unit": "Shares"},
-                        },
-                        {
-                            "date": "2024-04-12",
-                            "open": {"value": "419.23", "unit": "USD"},
-                            "high": {"value": "421.94", "unit": "USD"},
-                            "low": {"value": "416.45", "unit": "USD"},
-                            "close": {"value": "417.81", "unit": "USD"},
-                            "volume": {"value": "1182229", "unit": "Shares"},
-                        },
-                    ]
-                }
-            },
-            "errors": [
-                "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
-            ],
-        }
+        expected_response = GetPricesFromIdentifiersResp.model_validate(
+            {
+                "results": {
+                    "SPGI": {
+                        "prices": [
+                            {
+                                "date": "2024-04-11",
+                                "open": {"value": "424.26", "unit": "USD"},
+                                "high": {"value": "425.99", "unit": "USD"},
+                                "low": {"value": "422.04", "unit": "USD"},
+                                "close": {"value": "422.92", "unit": "USD"},
+                                "volume": {"value": "1129158", "unit": "Shares"},
+                            },
+                            {
+                                "date": "2024-04-12",
+                                "open": {"value": "419.23", "unit": "USD"},
+                                "high": {"value": "421.94", "unit": "USD"},
+                                "low": {"value": "416.45", "unit": "USD"},
+                                "close": {"value": "417.81", "unit": "USD"},
+                                "volume": {"value": "1182229", "unit": "Shares"},
+                            },
+                        ]
+                    }
+                },
+                "errors": [
+                    "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
+                ],
+            }
+        )
 
         tool = GetPricesFromIdentifiers(kfinance_client=mock_client)
         response = tool.run(
@@ -142,12 +148,14 @@ class TestGetPricesFromIdentifiers:
                 }
             ]
         }
-        expected_response = {
-            "results": {
-                "C_1": expected_single_company_response,
-                "C_2": expected_single_company_response,
-            },
-        }
+        expected_response = GetPricesFromIdentifiersResp.model_validate(
+            {
+                "results": {
+                    "C_1": expected_single_company_response,
+                    "C_2": expected_single_company_response,
+                },
+            }
+        )
         tool = GetPricesFromIdentifiers(kfinance_client=mock_client)
         response = tool.run(
             GetPricesFromIdentifiersArgs(

kfinance/domains/segments/segment_tools.py

@@ -43,7 +43,7 @@ class GetSegmentsFromIdentifiers(KfinanceTool):
         end_year: int | None = None,
         start_quarter: Literal[1, 2, 3, 4] | None = None,
         end_quarter: Literal[1, 2, 3, 4] | None = None,
-    ) -> dict:
+    ) -> GetSegmentsFromIdentifiersResp:
         """Sample Response:
 
         {
@@ -102,7 +102,6 @@ class GetSegmentsFromIdentifiers(KfinanceTool):
                 most_recent_year_data = segments_response.segments[most_recent_year]
                 segments_response.segments = {most_recent_year: most_recent_year_data}
 
-        output_model = GetSegmentsFromIdentifiersResp(
+        return GetSegmentsFromIdentifiersResp(
             results=segments_responses, errors=list(id_triple_resp.errors.values())
         )
-        return output_model.model_dump(mode="json")

kfinance/domains/segments/tests/test_segment_tools.py

@@ -7,6 +7,7 @@ from kfinance.domains.segments.segment_models import SegmentType
 from kfinance.domains.segments.segment_tools import (
     GetSegmentsFromIdentifiers,
     GetSegmentsFromIdentifiersArgs,
+    GetSegmentsFromIdentifiersResp,
 )
 
 
@@ -43,12 +44,14 @@ class TestGetSegmentsFromIdentifier:
             json=self.segments_response,
         )
 
-        expected_response = {
-            "results": {"SPGI": self.segments_response},
-            "errors": [
-                "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
-            ],
-        }
+        expected_response = GetSegmentsFromIdentifiersResp.model_validate(
+            {
+                "results": {"SPGI": self.segments_response},
+                "errors": [
+                    "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
+                ],
+            }
+        )
 
         tool = GetSegmentsFromIdentifiers(kfinance_client=mock_client)
         args = GetSegmentsFromIdentifiersArgs(
@@ -65,12 +68,14 @@ class TestGetSegmentsFromIdentifier:
         """
 
         company_ids = [1, 2]
-        expected_response = {
-            "results": {
-                "C_1": {"segments": {"2021": self.segments_response["segments"]["2021"]}},
-                "C_2": {"segments": {"2021": self.segments_response["segments"]["2021"]}},
+        expected_response = GetSegmentsFromIdentifiersResp.model_validate(
+            {
+                "results": {
+                    "C_1": {"segments": {"2021": self.segments_response["segments"]["2021"]}},
+                    "C_2": {"segments": {"2021": self.segments_response["segments"]["2021"]}},
+                }
             }
-        }
+        )
 
         for company_id in company_ids:
             requests_mock.get(

kfinance/domains/statements/statement_tools.py

@@ -55,7 +55,7 @@ class GetFinancialStatementFromIdentifiers(KfinanceTool):
         end_year: int | None = None,
         start_quarter: Literal[1, 2, 3, 4] | None = None,
         end_quarter: Literal[1, 2, 3, 4] | None = None,
-    ) -> dict:
+    ) -> GetFinancialStatementFromIdentifiersResp:
         """Sample response:
 
         {
@@ -109,7 +109,6 @@ class GetFinancialStatementFromIdentifiers(KfinanceTool):
                 most_recent_year_data = statement_response.statements[most_recent_year]
                 statement_response.statements = {most_recent_year: most_recent_year_data}
 
-        output_model = GetFinancialStatementFromIdentifiersResp(
+        return GetFinancialStatementFromIdentifiersResp(
             results=statement_responses, errors=list(id_triple_resp.errors.values())
         )
-        return output_model.model_dump(mode="json")

kfinance/domains/statements/tests/test_statement_tools.py

@@ -7,6 +7,7 @@ from kfinance.domains.statements.statement_models import StatementType
 from kfinance.domains.statements.statement_tools import (
     GetFinancialStatementFromIdentifiers,
     GetFinancialStatementFromIdentifiersArgs,
+    GetFinancialStatementFromIdentifiersResp,
 )
 
 
@@ -31,25 +32,27 @@ class TestGetFinancialStatementFromIdentifiers:
             url=f"https://kfinance.kensho.com/api/v1/statements/{SPGI_COMPANY_ID}/income_statement/none/none/none/none/none",
             json=self.statement_resp,
         )
-        expected_response = {
-            "results": {
-                "SPGI": {
-                    "statements": {
-                        "2020": {
-                            "Revenues": "7442000000.000000",
-                            "Total Revenues": "7442000000.000000",
-                        },
-                        "2021": {
-                            "Revenues": "8243000000.000000",
-                            "Total Revenues": "8243000000.000000",
-                        },
+        expected_response = GetFinancialStatementFromIdentifiersResp.model_validate(
+            {
+                "results": {
+                    "SPGI": {
+                        "statements": {
+                            "2020": {
+                                "Revenues": "7442000000.000000",
+                                "Total Revenues": "7442000000.000000",
+                            },
+                            "2021": {
+                                "Revenues": "8243000000.000000",
+                                "Total Revenues": "8243000000.000000",
+                            },
+                        }
                     }
-                }
-            },
-            "errors": [
-                "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
-            ],
-        }
+                },
+                "errors": [
+                    "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
+                ],
+            }
+        )
 
         tool = GetFinancialStatementFromIdentifiers(kfinance_client=mock_client)
         args = GetFinancialStatementFromIdentifiersArgs(
@@ -66,26 +69,28 @@ class TestGetFinancialStatementFromIdentifiers:
         """
 
         company_ids = [1, 2]
-        expected_response = {
-            "results": {
-                "C_1": {
-                    "statements": {
-                        "2021": {
-                            "Revenues": "8243000000.000000",
-                            "Total Revenues": "8243000000.000000",
+        expected_response = GetFinancialStatementFromIdentifiersResp.model_validate(
+            {
+                "results": {
+                    "C_1": {
+                        "statements": {
+                            "2021": {
+                                "Revenues": "8243000000.000000",
+                                "Total Revenues": "8243000000.000000",
+                            }
                         }
-                    }
-                },
-                "C_2": {
-                    "statements": {
-                        "2021": {
-                            "Revenues": "8243000000.000000",
-                            "Total Revenues": "8243000000.000000",
+                    },
+                    "C_2": {
+                        "statements": {
+                            "2021": {
+                                "Revenues": "8243000000.000000",
+                                "Total Revenues": "8243000000.000000",
+                            }
                         }
-                    }
-                },
+                    },
+                }
             }
-        }
+        )
 
         for company_id in company_ids:
             requests_mock.get(

kfinance/integrations/tool_calling/static_tools/get_n_quarters_ago.py

@@ -11,6 +11,11 @@ class GetNQuartersAgoArgs(BaseModel):
     n: int = Field(description="Number of quarters before the current quarter")
 
 
+class GetNQuartersAgoResp(BaseModel):
+    year: int
+    quarter: int
+
+
 class GetNQuartersAgo(KfinanceTool):
     name: str = "get_n_quarters_ago"
     description: str = (

kfinance/integrations/tool_calling/static_tools/tests/test_get_lastest.py

@@ -3,6 +3,7 @@ from datetime import datetime
 import time_machine
 
 from kfinance.client.kfinance import Client
+from kfinance.client.models.date_and_period_models import LatestPeriods
 from kfinance.integrations.tool_calling.static_tools.get_latest import GetLatest, GetLatestArgs
 
 
@@ -15,16 +16,18 @@ class TestGetLatest:
         THEN we get back latest info
         """
 
-        expected_resp = {
-            "annual": {"latest_year": 2024},
-            "now": {
-                "current_date": "2025-01-01",
-                "current_month": 1,
-                "current_quarter": 1,
-                "current_year": 2025,
-            },
-            "quarterly": {"latest_quarter": 4, "latest_year": 2024},
-        }
+        expected_resp = LatestPeriods.model_validate(
+            {
+                "annual": {"latest_year": 2024},
+                "now": {
+                    "current_date": "2025-01-01",
+                    "current_month": 1,
+                    "current_quarter": 1,
+                    "current_year": 2025,
+                },
+                "quarterly": {"latest_quarter": 4, "latest_year": 2024},
+            }
+        )
         tool = GetLatest(kfinance_client=mock_client)
         resp = tool.run(GetLatestArgs().model_dump(mode="json"))
         assert resp == expected_resp

kfinance/integrations/tool_calling/static_tools/tests/test_get_n_quarters_ago.py

@@ -3,6 +3,7 @@ from datetime import datetime
 import time_machine
 
 from kfinance.client.kfinance import Client
+from kfinance.client.models.date_and_period_models import YearAndQuarter
 from kfinance.integrations.tool_calling.static_tools.get_n_quarters_ago import (
     GetNQuartersAgo,
     GetNQuartersAgoArgs,
@@ -18,7 +19,7 @@ class TestGetNQuartersAgo:
         THEN we get back 3 quarters ago
         """
 
-        expected_resp = {"quarter": 2, "year": 2024}
+        expected_resp = YearAndQuarter(year=2024, quarter=2)
         tool = GetNQuartersAgo(kfinance_client=mock_client)
         resp = tool.run(GetNQuartersAgoArgs(n=3).model_dump(mode="json"))
         assert resp == expected_resp

kfinance/integrations/tool_calling/tests/test_tool_calling_models.py

@@ -7,7 +7,10 @@ from requests_mock import Mocker
 
 from kfinance.client.kfinance import Client
 from kfinance.conftest import SPGI_COMPANY_ID
-from kfinance.domains.companies.company_tools import GetInfoFromIdentifiers
+from kfinance.domains.companies.company_tools import (
+    GetInfoFromIdentifiers,
+    GetInfoFromIdentifiersResp,
+)
 from kfinance.integrations.tool_calling.tool_calling_models import ValidQuarter
 
 
@@ -27,7 +30,10 @@ class TestGetEndpointsFromToolCallsWithGrounding:
             "https://kfinance.kensho.com/api/v1/ids",
             "https://kfinance.kensho.com/api/v1/info/21719",
         ]
-        expected_resp = {"data": {"results": {"SPGI": resp_data}}, "endpoint_urls": resp_endpoint}
+        expected_resp = {
+            "data": GetInfoFromIdentifiersResp.model_validate({"results": {"SPGI": resp_data}}),
+            "endpoint_urls": resp_endpoint,
+        }
 
         requests_mock.get(
             url=f"https://kfinance.kensho.com/api/v1/info/{SPGI_COMPANY_ID}",

kfinance/integrations/tool_calling/tool_calling_models.py

@@ -1,3 +1,4 @@
+import abc
 from typing import Annotated, Any, Callable, Dict, Literal, Type
 
 from langchain_core.tools import BaseTool
@@ -22,7 +23,7 @@ class KfinanceTool(BaseTool):
 
     model_config = ConfigDict(extra="forbid")
 
-    def run_without_langchain(self, *args: Any, **kwargs: Any) -> Any:
+    def run_without_langchain(self, *args: Any, **kwargs: Any) -> dict:
         """Execute a Kfinance tool without langchain.
 
         Langchain converts json input params into the pydantic args_schema, which means that
@@ -38,7 +39,8 @@ class KfinanceTool(BaseTool):
         # This behavior matches the langchain handling. See
         # https://github.com/langchain-ai/langchain/blob/ca39680d2ab0d786bc035930778a5787e7bb5e01/libs/core/langchain_core/tools/base.py#L595-L597
         args_dict = {k: v for k, v in args_dict.items() if k in kwargs}
-        return self._run(**args_dict)
+        result_model = self._run(**args_dict)
+        return result_model.model_dump(mode="json", exclude_none=True)
 
     def run_with_grounding(self, *args: Any, **kwargs: Any) -> Any:
         """Execute a Kfinance tool with grounding support.
@@ -47,7 +49,10 @@ class KfinanceTool(BaseTool):
         support, for returning the endpoint urls along with the data as citation info for the LRA Data Agent.
         """
         with self.kfinance_client.kfinance_api_client.endpoint_tracker() as endpoint_tracker_queue:
-            data = self.run_without_langchain(*args, **kwargs)
+            args_model = self.args_schema.model_validate(kwargs)
+            args_dict = args_model.model_dump()
+            args_dict = {k: v for k, v in args_dict.items() if k in kwargs}
+            result_model = self._run(**args_dict)
 
             # After completion of tool data fetching and within the endpoint_tracker context manager scope, dequeue the endpoint_tracker_queue
             endpoint_urls = []
@@ -55,11 +60,12 @@ class KfinanceTool(BaseTool):
                 endpoint_urls.append(endpoint_tracker_queue.get())
 
             return {
-                "data": data,
+                "data": result_model,
                 "endpoint_urls": endpoint_urls,
             }
 
-    def _run(self, *args: Any, **kwargs: Any) -> Any:
+    @abc.abstractmethod
+    def _run(self, *args: Any, **kwargs: Any) -> BaseModel:
         """The code to execute the tool.
 
         Where feasible and useful, tools should use batch processing to parallelize

kfinance/version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '3.2.5'
-__version_tuple__ = version_tuple = (3, 2, 5)
+__version__ = version = '3.2.7'
+__version_tuple__ = version_tuple = (3, 2, 7)
 
 __commit_id__ = commit_id = None