kensho-kfinance 3.2.4__py3-none-any.whl → 4.0.0__py3-none-any.whl

kfinance/domains/statements/tests/test_statement_tools.py

@@ -1,21 +1,46 @@
 from requests_mock import Mocker
 
 from kfinance.client.kfinance import Client
-from kfinance.conftest import SPGI_COMPANY_ID
 from kfinance.domains.companies.company_models import COMPANY_ID_PREFIX
 from kfinance.domains.statements.statement_models import StatementType
 from kfinance.domains.statements.statement_tools import (
     GetFinancialStatementFromIdentifiers,
     GetFinancialStatementFromIdentifiersArgs,
+    GetFinancialStatementFromIdentifiersResp,
 )
 
 
 class TestGetFinancialStatementFromIdentifiers:
     statement_resp = {
-        "statements": {
-            "2020": {"Revenues": "7442000000.000000", "Total Revenues": "7442000000.000000"},
-            "2021": {"Revenues": "8243000000.000000", "Total Revenues": "8243000000.000000"},
-        }
+        "currency": "USD",
+        "periods": {
+            "CY2020": {
+                "period_end_date": "2020-12-31",
+                "num_months": 12,
+                "statements": [
+                    {
+                        "name": "Income Statement",
+                        "line_items": [
+                            {"name": "Revenues", "value": "7442000000.000000", "sources": []},
+                            {"name": "Total Revenues", "value": "7442000000.000000", "sources": []},
+                        ],
+                    }
+                ],
+            },
+            "CY2021": {
+                "period_end_date": "2021-12-31",
+                "num_months": 12,
+                "statements": [
+                    {
+                        "name": "Income Statement",
+                        "line_items": [
+                            {"name": "Revenues", "value": "8243000000.000000", "sources": []},
+                            {"name": "Total Revenues", "value": "8243000000.000000", "sources": []},
+                        ],
+                    }
+                ],
+            },
+        },
     }
 
     def test_get_financial_statement_from_identifiers(
@@ -27,33 +52,88 @@ class TestGetFinancialStatementFromIdentifiers:
         THEN we get back the SPGI income statement and an error for the non-existent company.
         """
 
-        requests_mock.get(
-            url=f"https://kfinance.kensho.com/api/v1/statements/{SPGI_COMPANY_ID}/income_statement/none/none/none/none/none",
-            json=self.statement_resp,
+        # Mock the unified_fetch_id_triples response
+        requests_mock.post(
+            url="https://kfinance.kensho.com/api/v1/ids",
+            json={
+                "identifiers_to_id_triples": {
+                    "SPGI": {
+                        "company_id": 21719,
+                        "security_id": 2629107,
+                        "trading_item_id": 2629108,
+                    }
+                },
+                "errors": {
+                    "NON-EXISTENT": "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
+                },
+            },
+        )
+
+        # Mock the fetch_statement response
+        requests_mock.post(
+            url="https://kfinance.kensho.com/api/v1/statements/",
+            json={"results": {"21719": self.statement_resp}, "errors": {}},
         )
-        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": {
+                        "currency": "USD",
+                        "periods": {
+                            "CY2020": {
+                                "period_end_date": "2020-12-31",
+                                "num_months": 12,
+                                "statements": [
+                                    {
+                                        "name": "Income Statement",
+                                        "line_items": [
+                                            {
+                                                "name": "Revenues",
+                                                "value": "7442000000.000000",
+                                                "sources": [],
+                                            },
+                                            {
+                                                "name": "Total Revenues",
+                                                "value": "7442000000.000000",
+                                                "sources": [],
+                                            },
+                                        ],
+                                    }
+                                ],
+                            },
+                            "CY2021": {
+                                "period_end_date": "2021-12-31",
+                                "num_months": 12,
+                                "statements": [
+                                    {
+                                        "name": "Income Statement",
+                                        "line_items": [
+                                            {
+                                                "name": "Revenues",
+                                                "value": "8243000000.000000",
+                                                "sources": [],
+                                            },
+                                            {
+                                                "name": "Total Revenues",
+                                                "value": "8243000000.000000",
+                                                "sources": [],
+                                            },
+                                        ],
+                                    }
+                                ],
+                            },
                         },
                     }
-                }
-            },
-            "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(
-            identifiers=["SPGI", "non-existent"], statement=StatementType.income_statement
+            identifiers=["SPGI", "NON-EXISTENT"], statement=StatementType.income_statement
         )
         response = tool.run(args.model_dump(mode="json"))
         assert response == expected_response
@@ -66,32 +146,154 @@ class TestGetFinancialStatementFromIdentifiers:
         """
 
         company_ids = [1, 2]
-        expected_response = {
-            "results": {
-                "C_1": {
-                    "statements": {
-                        "2021": {
-                            "Revenues": "8243000000.000000",
-                            "Total Revenues": "8243000000.000000",
-                        }
-                    }
-                },
-                "C_2": {
-                    "statements": {
-                        "2021": {
-                            "Revenues": "8243000000.000000",
-                            "Total Revenues": "8243000000.000000",
-                        }
-                    }
+        expected_response = GetFinancialStatementFromIdentifiersResp.model_validate(
+            {
+                "results": {
+                    "C_1": {
+                        "currency": "USD",
+                        "periods": {
+                            "CY2021": {
+                                "period_end_date": "2021-12-31",
+                                "num_months": 12,
+                                "statements": [
+                                    {
+                                        "name": "Income Statement",
+                                        "line_items": [
+                                            {
+                                                "name": "Revenues",
+                                                "value": "8243000000.000000",
+                                                "sources": [],
+                                            },
+                                            {
+                                                "name": "Total Revenues",
+                                                "value": "8243000000.000000",
+                                                "sources": [],
+                                            },
+                                        ],
+                                    }
+                                ],
+                            }
+                        },
+                    },
+                    "C_2": {
+                        "currency": "USD",
+                        "periods": {
+                            "CY2021": {
+                                "period_end_date": "2021-12-31",
+                                "num_months": 12,
+                                "statements": [
+                                    {
+                                        "name": "Income Statement",
+                                        "line_items": [
+                                            {
+                                                "name": "Revenues",
+                                                "value": "8243000000.000000",
+                                                "sources": [],
+                                            },
+                                            {
+                                                "name": "Total Revenues",
+                                                "value": "8243000000.000000",
+                                                "sources": [],
+                                            },
+                                        ],
+                                    }
+                                ],
+                            }
+                        },
+                    },
+                }
+            }
+        )
+
+        # Mock the unified_fetch_id_triples response
+        requests_mock.post(
+            url="https://kfinance.kensho.com/api/v1/ids",
+            json={
+                "identifiers_to_id_triples": {
+                    "C_1": {"company_id": 1, "security_id": 101, "trading_item_id": 201},
+                    "C_2": {"company_id": 2, "security_id": 102, "trading_item_id": 202},
                 },
+                "errors": {},
+            },
+        )
+
+        # Mock the fetch_statement response
+        requests_mock.post(
+            url="https://kfinance.kensho.com/api/v1/statements/",
+            json={"results": {"1": self.statement_resp, "2": self.statement_resp}, "errors": {}},
+        )
+
+        tool = GetFinancialStatementFromIdentifiers(kfinance_client=mock_client)
+        args = GetFinancialStatementFromIdentifiersArgs(
+            identifiers=[f"{COMPANY_ID_PREFIX}{company_id}" for company_id in company_ids],
+            statement=StatementType.income_statement,
+        )
+        response = tool.run(args.model_dump(mode="json"))
+        assert response == expected_response
+
+    def test_empty_most_recent_request(self, requests_mock: Mocker, mock_client: Client) -> None:
+        """
+        GIVEN the GetFinancialStatementFromIdentifiers tool
+        WHEN we request most recent statement for multiple companies
+        THEN we only get back the most recent statement for each company
+        UNLESS no statements exist
+        """
+
+        company_ids = [1, 2]
+        expected_response = GetFinancialStatementFromIdentifiersResp.model_validate(
+            {
+                "results": {
+                    "C_1": {"currency": "USD", "periods": {}},
+                    "C_2": {
+                        "currency": "USD",
+                        "periods": {
+                            "CY2021": {
+                                "period_end_date": "2021-12-31",
+                                "num_months": 12,
+                                "statements": [
+                                    {
+                                        "name": "Income Statement",
+                                        "line_items": [
+                                            {
+                                                "name": "Revenues",
+                                                "value": "8243000000.000000",
+                                                "sources": [],
+                                            },
+                                            {
+                                                "name": "Total Revenues",
+                                                "value": "8243000000.000000",
+                                                "sources": [],
+                                            },
+                                        ],
+                                    }
+                                ],
+                            }
+                        },
+                    },
+                }
             }
-        }
+        )
 
-        for company_id in company_ids:
-            requests_mock.get(
-                url=f"https://kfinance.kensho.com/api/v1/statements/{company_id}/income_statement/none/none/none/none/none",
-                json=self.statement_resp,
-            )
+        # Mock the unified_fetch_id_triples response
+        requests_mock.post(
+            url="https://kfinance.kensho.com/api/v1/ids",
+            json={
+                "identifiers_to_id_triples": {
+                    "C_1": {"company_id": 1, "security_id": 101, "trading_item_id": 201},
+                    "C_2": {"company_id": 2, "security_id": 102, "trading_item_id": 202},
+                },
+                "errors": {},
+            },
+        )
+
+        # Mock the fetch_statement response with different data for different companies
+        requests_mock.post(
+            url="https://kfinance.kensho.com/api/v1/statements/",
+            json={
+                "results": {"1": {"currency": "USD", "periods": {}}, "2": self.statement_resp},
+                "errors": {},
+            },
+        )
 
         tool = GetFinancialStatementFromIdentifiers(kfinance_client=mock_client)
         args = GetFinancialStatementFromIdentifiersArgs(

kfinance/integrations/local_mcp/kfinance_mcp.py

@@ -19,4 +19,4 @@ class KfinanceMcp(FastMCP):
         call_tool(validate_input=False) turns off the mcp sdk validation.
         """
         super()._setup_handlers()
-        self._mcp_server.call_tool(validate_input=False)(self._mcp_call_tool)
+        self._mcp_server.call_tool(validate_input=False)(self._call_tool_mcp)

kfinance/integrations/tests/test_example_notebook.py

@@ -55,6 +55,7 @@ def test_run_notebook(jupyter_kernel_name: str):
     # - mocks for all calls made by the client while executing the notebook
     startup_cell_code = dedent("""
         from datetime import datetime
+        from urllib.parse import quote
         from kfinance.client.kfinance import Client
         kfinance_client = Client(refresh_token="foo")
         api_client = kfinance_client.kfinance_api_client
@@ -85,22 +86,44 @@ def test_run_notebook(jupyter_kernel_name: str):
         )
 
         balance_sheet_resp = {
-            "statements": {
-                "2022Q3": {"Cash And Equivalents": "1387000000.000000"},
-                "2022Q4": {"Cash And Equivalents": "1286000000.000000"}
+            "currency": "USD",
+            "periods": {
+                "CY2022Q3": {
+                    "period_end_date": "2022-09-30",
+                    "num_months": 3,
+                    "statements": [
+                        {
+                            "name": "Balance Sheet",
+                            "line_items": [
+                                {"name": "Cash And Equivalents", "value": "1387000000.000000", "sources": []}
+                            ]
+                        }
+                    ]
+                },
+                "CY2022Q4": {
+                    "period_end_date": "2022-12-31",
+                    "num_months": 3,
+                    "statements": [
+                        {
+                            "name": "Balance Sheet",
+                            "line_items": [
+                                {"name": "Cash And Equivalents", "value": "1286000000.000000", "sources": []}
+                            ]
+                        }
+                    ]
+                }
             }
         }
 
-        # spgi.balance_sheet()
-        mocker.get(
-            url="https://kfinance.kensho.com/api/v1/statements/21719/balance_sheet/none/none/none/none/none",
-            json=balance_sheet_resp
-        )
-
-        # spgi.balance_sheet(period_type=PeriodType.annual, start_year=2010, end_year=2019)
-        mocker.get(
-            url="https://kfinance.kensho.com/api/v1/statements/21719/balance_sheet/annual/2010/2019/none/none",
-            json=balance_sheet_resp
+        # spgi.balance_sheet() and spgi.balance_sheet(period_type=PeriodType.annual, start_year=2010, end_year=2019)
+        mocker.post(
+            url="https://kfinance.kensho.com/api/v1/statements/",
+            json={
+                "results": {
+                    "21719": balance_sheet_resp
+                },
+                "errors": {}
+            }
         )
 
         # kfinance_client.ticker("JPM").balance_sheet()
@@ -112,9 +135,27 @@ def test_run_notebook(jupyter_kernel_name: str):
         )
 
         # spgi.net_income(period_type=PeriodType.annual, start_year=2010, end_year=2019)
-        mocker.get(
-            url="https://kfinance.kensho.com/api/v1/line_item/21719/net_income/annual/2010/2019/none/none",
-            json={"line_item": {"2010": "828000000.000000"}}
+        mocker.post(
+            url="https://kfinance.kensho.com/api/v1/line_item/",
+            json={
+                "results": {
+                    "21719": {
+                        "currency": "USD",
+                        "periods": {
+                            "CY2010": {
+                                "period_end_date": "2010-12-31",
+                                "num_months": 12,
+                                "line_item": {
+                                    "name": "Net Income",
+                                    "value": "828000000.000000",
+                                    "sources": []
+                                }
+                            }
+                        }
+                    }
+                },
+                "errors": {}
+            }
         )
 
         prices_resp = {

kfinance/integrations/tool_calling/all_tools.py

@@ -25,7 +25,10 @@ from kfinance.domains.mergers_and_acquisitions.merger_and_acquisition_tools impo
     GetMergerInfoFromTransactionId,
     GetMergersFromIdentifiers,
 )
-from kfinance.domains.prices.price_tools import GetPricesFromIdentifiers
+from kfinance.domains.prices.price_tools import (
+    GetHistoryMetadataFromIdentifiers,
+    GetPricesFromIdentifiers,
+)
 from kfinance.domains.segments.segment_tools import GetSegmentsFromIdentifiers
 from kfinance.domains.statements.statement_tools import GetFinancialStatementFromIdentifiers
 from kfinance.integrations.tool_calling.static_tools.get_latest import GetLatest
@@ -61,6 +64,7 @@ ALL_TOOLS: list[type[KfinanceTool]] = [
     GetFinancialLineItemFromIdentifiers,
     # Prices
     GetPricesFromIdentifiers,
+    GetHistoryMetadataFromIdentifiers,
     # Segments
     GetSegmentsFromIdentifiers,
     # Statements

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,11 @@ 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_models import COMPANY_ID_PREFIX
+from kfinance.domains.companies.company_tools import (
+    GetInfoFromIdentifiers,
+    GetInfoFromIdentifiersResp,
+)
 from kfinance.integrations.tool_calling.tool_calling_models import ValidQuarter
 
 
@@ -22,13 +26,20 @@ class TestGetEndpointsFromToolCallsWithGrounding:
         """
 
         # truncated from the original
-        resp_data = {"name": "S&P Global Inc.", "status": "Operating"}
+        resp_data = {
+            "name": "S&P Global Inc.",
+            "status": "Operating",
+            "company_id": f"{COMPANY_ID_PREFIX}{SPGI_COMPANY_ID}",
+        }
         resp_endpoint = [
             "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,
+        }
+        del resp_data["company_id"]
         requests_mock.get(
             url=f"https://kfinance.kensho.com/api/v1/info/{SPGI_COMPANY_ID}",
             json=resp_data,

kfinance/integrations/tool_calling/tool_calling_models.py

@@ -1,7 +1,14 @@
+import abc
 from typing import Annotated, Any, Callable, Dict, Literal, Type
 
 from langchain_core.tools import BaseTool
-from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, model_serializer
+from pydantic import (
+    BaseModel,
+    BeforeValidator,
+    ConfigDict,
+    Field,
+    model_serializer,
+)
 
 from kfinance.client.kfinance import Client
 from kfinance.client.permission_models import Permission
@@ -22,7 +29,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 +45,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 +55,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 +66,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.4'
-__version_tuple__ = version_tuple = (3, 2, 4)
+__version__ = version = '4.0.0'
+__version_tuple__ = version_tuple = (4, 0, 0)
 
 __commit_id__ = commit_id = None