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

kfinance/domains/mergers_and_acquisitions/tests/test_merger_and_acquisition_tools.py

@@ -8,13 +8,18 @@ from kfinance.client.tests.test_objects import (
     ordered,
 )
 from kfinance.conftest import SPGI_COMPANY_ID
-from kfinance.domains.companies.company_models import COMPANY_ID_PREFIX
+from kfinance.domains.mergers_and_acquisitions.merger_and_acquisition_models import (
+    AdvisorResp,
+    MergerInfo,
+)
 from kfinance.domains.mergers_and_acquisitions.merger_and_acquisition_tools import (
     GetAdvisorsForCompanyInTransactionFromIdentifier,
     GetAdvisorsForCompanyInTransactionFromIdentifierArgs,
+    GetAdvisorsForCompanyInTransactionFromIdentifierResp,
     GetMergerInfoFromTransactionId,
     GetMergerInfoFromTransactionIdArgs,
     GetMergersFromIdentifiers,
+    GetMergersFromIdentifiersResp,
 )
 from kfinance.integrations.tool_calling.tool_calling_models import ToolArgsWithIdentifiers
 
@@ -27,12 +32,14 @@ class TestGetMergersFromIdentifiers:
         THEN we get back the SPGI mergers and an error for the non-existent company"""
 
         merger_data = MERGERS_RESP.model_dump(mode="json")
-        expected_response = {
-            "results": {"SPGI": merger_data},
-            "errors": [
-                "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
-            ],
-        }
+        expected_response = GetMergersFromIdentifiersResp.model_validate(
+            {
+                "results": {"SPGI": merger_data},
+                "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/mergers/{SPGI_COMPANY_ID}", json=merger_data
         )
@@ -51,13 +58,21 @@ class TestGetCompaniesAdvisingCompanyInTransactionFromIdentifier:
             "advisor_company_name": "Kensho Technologies, Inc.",
             "advisor_type_name": "Professional Mongo Enjoyer",
         }
-        api_response = {"advisors": [deepcopy(advisor_data)]}
-        expected_response = {"results": [deepcopy(advisor_data)]}
-        expected_response["results"][0]["advisor_company_id"] = f"{COMPANY_ID_PREFIX}251994106"
+        expected_response = GetAdvisorsForCompanyInTransactionFromIdentifierResp(
+            results=[
+                AdvisorResp(
+                    advisor_company_id=251994106,
+                    advisor_company_name="Kensho Technologies, Inc.",
+                    advisor_type_name="Professional Mongo Enjoyer",
+                )
+            ],
+            errors=[],
+        )
+
         transaction_id = 554979212
         requests_mock.get(
             url=f"https://kfinance.kensho.com/api/v1/merger/info/{transaction_id}/advisors/{SPGI_COMPANY_ID}",
-            json=api_response,
+            json={"advisors": [deepcopy(advisor_data)]},
         )
         tool = GetAdvisorsForCompanyInTransactionFromIdentifier(kfinance_client=mock_client)
         args = GetAdvisorsForCompanyInTransactionFromIdentifierArgs(
@@ -67,14 +82,14 @@ class TestGetCompaniesAdvisingCompanyInTransactionFromIdentifier:
         assert response == expected_response
 
     def test_get_companies_advising_company_in_transaction_from_bad_identifier(
-        self, requests_mock: Mocker, mock_client: Client
+        self, mock_client: Client
     ):
-        expected_response = {
-            "results": [],
-            "errors": [
+        expected_response = GetAdvisorsForCompanyInTransactionFromIdentifierResp(
+            results=[],
+            errors=[
                 "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
             ],
-        }
+        )
         transaction_id = 554979212
         tool = GetAdvisorsForCompanyInTransactionFromIdentifier(kfinance_client=mock_client)
         args = GetAdvisorsForCompanyInTransactionFromIdentifierArgs(
@@ -86,30 +101,11 @@ class TestGetCompaniesAdvisingCompanyInTransactionFromIdentifier:
 
 class TestGetMergerInfoFromTransactionId:
     def test_get_merger_info_from_transaction_id(self, requests_mock: Mocker, mock_client: Client):
-        response_base = {
-            "timeline": [
-                {"status": "Announced", "date": "2000-09-12"},
-                {"status": "Closed", "date": "2000-09-12"},
-            ],
-            "participants": {},
-            "consideration": {
-                "currency_name": "US Dollar",
-                "current_calculated_gross_total_transaction_value": "51609375.000000",
-                "current_calculated_implied_equity_value": "51609375.000000",
-                "current_calculated_implied_enterprise_value": "51609375.000000",
-                "details": [
-                    {
-                        "scenario": "Stock Lump Sum",
-                        "subtype": "Common Equity",
-                        "cash_or_cash_equivalent_per_target_share_unit": None,
-                        "number_of_target_shares_sought": "1000000.000000",
-                        "current_calculated_gross_value_of_consideration": "51609375.000000",
-                    }
-                ],
-            },
-        }
-
-        participants_api_response = {
+        timeline_resp = [
+            {"status": "Announced", "date": "2000-09-12"},
+            {"status": "Closed", "date": "2000-09-12"},
+        ]
+        participants_resp = {
             "target": {"company_id": 31696, "company_name": "MongoMusic, Inc."},
             "buyers": [{"company_id": 21835, "company_name": "Microsoft Corporation"}],
             "sellers": [
@@ -117,33 +113,39 @@ class TestGetMergerInfoFromTransactionId:
                 {"company_id": 20087, "company_name": "Draper Richards, L.P."},
             ],
         }
-        api_response = deepcopy(response_base)
-        api_response["participants"] = participants_api_response
 
-        # Returned company IDs should be prefixed
-        expected_participants_tool_response = {
-            "target": {
-                "company_id": f"{COMPANY_ID_PREFIX}31696",
-                "company_name": "MongoMusic, Inc.",
-            },
-            "buyers": [
-                {"company_id": f"{COMPANY_ID_PREFIX}21835", "company_name": "Microsoft Corporation"}
-            ],
-            "sellers": [
-                {"company_id": f"{COMPANY_ID_PREFIX}18805", "company_name": "Angel Investors L.P."},
+        consideration_resp = {
+            "currency_name": "US Dollar",
+            "current_calculated_gross_total_transaction_value": "51609375.000000",
+            "current_calculated_implied_equity_value": "51609375.000000",
+            "current_calculated_implied_enterprise_value": "51609375.000000",
+            "details": [
                 {
-                    "company_id": f"{COMPANY_ID_PREFIX}20087",
-                    "company_name": "Draper Richards, L.P.",
-                },
+                    "scenario": "Stock Lump Sum",
+                    "subtype": "Common Equity",
+                    "cash_or_cash_equivalent_per_target_share_unit": None,
+                    "number_of_target_shares_sought": "1000000.000000",
+                    "current_calculated_gross_value_of_consideration": "51609375.000000",
+                }
             ],
         }
-        expected_response = deepcopy(response_base)
-        expected_response["participants"] = expected_participants_tool_response
+
+        expected_response = MergerInfo.model_validate(
+            {
+                "timeline": timeline_resp,
+                "participants": participants_resp,
+                "consideration": consideration_resp,
+            }
+        )
 
         transaction_id = 517414
         requests_mock.get(
             url=f"https://kfinance.kensho.com/api/v1/merger/info/{transaction_id}",
-            json=api_response,
+            json={
+                "timeline": timeline_resp,
+                "participants": participants_resp,
+                "consideration": consideration_resp,
+            },
         )
         tool = GetMergerInfoFromTransactionId(kfinance_client=mock_client)
         args = GetMergerInfoFromTransactionIdArgs(transaction_id=transaction_id)

kfinance/domains/prices/price_tools.py

@@ -60,7 +60,7 @@ class GetPricesFromIdentifiers(KfinanceTool):
         end_date: date | None = None,
         periodicity: Periodicity = Periodicity.day,
         adjusted: bool = True,
-    ) -> dict:
+    ) -> GetPricesFromIdentifiersResp:
         """Sample Response:
 
         {
@@ -116,10 +116,9 @@ class GetPricesFromIdentifiers(KfinanceTool):
             for price_response in price_responses.values():
                 price_response.prices = price_response.prices[-1:]
 
-        output_model = GetPricesFromIdentifiersResp(
+        return GetPricesFromIdentifiersResp(
             results=price_responses, errors=list(id_triple_resp.errors.values())
         )
-        return output_model.model_dump(mode="json")
 
 
 class GetHistoryMetadataFromIdentifiersResp(ToolRespWithErrors):
@@ -136,7 +135,7 @@ class GetHistoryMetadataFromIdentifiers(KfinanceTool):
     args_schema: Type[BaseModel] = ToolArgsWithIdentifiers
     accepted_permissions: set[Permission] | None = None
 
-    def _run(self, identifiers: list[str]) -> dict:
+    def _run(self, identifiers: list[str]) -> GetHistoryMetadataFromIdentifiersResp:
         """Sample response:
 
         {
@@ -169,8 +168,6 @@ class GetHistoryMetadataFromIdentifiers(KfinanceTool):
         history_metadata_responses: dict[str, HistoryMetadataResp] = (
             process_tasks_in_thread_pool_executor(api_client=api_client, tasks=tasks)
         )
-        output_model = GetHistoryMetadataFromIdentifiersResp(
+        return GetHistoryMetadataFromIdentifiersResp(
             results=history_metadata_responses, errors=list(id_triple_resp.errors.values())
         )
-
-        return output_model.model_dump(mode="json")

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}",