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

kfinance/domains/business_relationships/tests/test_business_relationship_tools.py

@@ -8,6 +8,7 @@ from kfinance.domains.business_relationships.business_relationship_models import
 from kfinance.domains.business_relationships.business_relationship_tools import (
     GetBusinessRelationshipFromIdentifiers,
     GetBusinessRelationshipFromIdentifiersArgs,
+    GetBusinessRelationshipFromIdentifiersResp,
 )
 
 
@@ -27,21 +28,23 @@ class TestGetBusinessRelationshipFromIdentifiers:
                 {"company_id": 8182358, "company_name": "Eloqua, Inc."},
             ],
         }
-        expected_result = {
-            "business_relationship": "supplier",
-            "results": {
-                "SPGI": {
-                    "current": [{"company_id": "C_883103", "company_name": "CRISIL Limited"}],
-                    "previous": [
-                        {"company_id": "C_472898", "company_name": "Morgan Stanley"},
-                        {"company_id": "C_8182358", "company_name": "Eloqua, Inc."},
-                    ],
-                }
-            },
-            "errors": [
-                "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
-            ],
-        }
+        expected_result = GetBusinessRelationshipFromIdentifiersResp.model_validate(
+            {
+                "business_relationship": "supplier",
+                "results": {
+                    "SPGI": {
+                        "current": [{"company_id": 883103, "company_name": "CRISIL Limited"}],
+                        "previous": [
+                            {"company_id": 472898, "company_name": "Morgan Stanley"},
+                            {"company_id": 8182358, "company_name": "Eloqua, Inc."},
+                        ],
+                    }
+                },
+                "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/relationship/{SPGI_COMPANY_ID}/supplier",
@@ -54,5 +57,4 @@ class TestGetBusinessRelationshipFromIdentifiers:
             business_relationship=BusinessRelationshipType.supplier,
         )
         resp = tool.run(args.model_dump(mode="json"))
-        resp["results"]["SPGI"]["previous"].sort(key=lambda x: x["company_id"])
         assert resp == expected_result

kfinance/domains/capitalizations/capitalization_models.py

@@ -20,9 +20,9 @@ class DailyCapitalization(BaseModel):
     """DailyCapitalization represents market cap, TEV, and shares outstanding for a day"""
 
     date: date
-    market_cap: Money
-    tev: Money
-    shares_outstanding: Shares
+    market_cap: Money | None
+    tev: Money | None
+    shares_outstanding: Shares | None
 
 
 class Capitalizations(BaseModel):

kfinance/domains/capitalizations/capitalization_tools.py

@@ -26,6 +26,7 @@ class GetCapitalizationFromIdentifiersArgs(ToolArgsWithIdentifiers):
 
 
 class GetCapitalizationFromIdentifiersResp(ToolRespWithErrors):
+    capitalization: Capitalization
     results: dict[str, Capitalizations]
 
 
@@ -50,13 +51,13 @@ class GetCapitalizationFromIdentifiers(KfinanceTool):
         capitalization: Capitalization,
         start_date: str | None = None,
         end_date: str | None = None,
-    ) -> dict:
+    ) -> GetCapitalizationFromIdentifiersResp:
         """Sample response:
 
         {
+            'capitalization': 'market_cap'
             'results': {
                 'SPGI': {
-                    'capitalizations': [
                         {'date': '2024-04-10', 'market_cap': {'value': '132766738270.00', 'unit': 'USD'}},
                         {'date': '2024-04-11', 'market_cap': {'value': '132416066761.00', 'unit': 'USD'}}
                     ]
@@ -100,7 +101,8 @@ class GetCapitalizationFromIdentifiers(KfinanceTool):
                 if capitalization is not Capitalization.shares_outstanding:
                     daily_capitalization.shares_outstanding = None
 
-        resp_model = GetCapitalizationFromIdentifiersResp(
-            results=capitalization_responses, errors=list(id_triple_resp.errors.values())
+        return GetCapitalizationFromIdentifiersResp(
+            capitalization=capitalization,
+            results=capitalization_responses,
+            errors=list(id_triple_resp.errors.values()),
         )
-        return resp_model.model_dump(mode="json", exclude_none=True)

kfinance/domains/capitalizations/tests/test_capitalization_tools.py

@@ -1,11 +1,20 @@
+from datetime import date
+from decimal import Decimal
+
 from requests_mock import Mocker
 
 from kfinance.client.kfinance import Client
+from kfinance.client.models.decimal_with_unit import Money
 from kfinance.conftest import SPGI_COMPANY_ID
-from kfinance.domains.capitalizations.capitalization_models import Capitalization
+from kfinance.domains.capitalizations.capitalization_models import (
+    Capitalization,
+    Capitalizations,
+    DailyCapitalization,
+)
 from kfinance.domains.capitalizations.capitalization_tools import (
     GetCapitalizationFromIdentifiers,
     GetCapitalizationFromIdentifiersArgs,
+    GetCapitalizationFromIdentifiersResp,
 )
 from kfinance.domains.companies.company_models import COMPANY_ID_PREFIX
 
@@ -40,25 +49,30 @@ class TestGetCapitalizationFromCompanyIds:
             json=self.market_caps_resp,
         )
 
-        expected_response = {
-            "results": {
-                "SPGI": {
-                    "capitalizations": [
-                        {
-                            "date": "2024-04-10",
-                            "market_cap": {"value": "132766738270.00", "unit": "USD"},
-                        },
-                        {
-                            "date": "2024-04-11",
-                            "market_cap": {"value": "132416066761.00", "unit": "USD"},
-                        },
+        expected_response = GetCapitalizationFromIdentifiersResp(
+            capitalization=Capitalization.market_cap,
+            results={
+                "SPGI": Capitalizations(
+                    market_caps=[
+                        DailyCapitalization(
+                            date=date(2024, 4, 10),
+                            market_cap=Money(value=Decimal(132766738270), unit="USD"),
+                            tev=None,
+                            shares_outstanding=None,
+                        ),
+                        DailyCapitalization(
+                            date=date(2024, 4, 11),
+                            market_cap=Money(value=Decimal(132416066761), unit="USD"),
+                            tev=None,
+                            shares_outstanding=None,
+                        ),
                     ]
-                }
+                )
             },
-            "errors": [
+            errors=[
                 "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
             ],
-        }
+        )
 
         tool = GetCapitalizationFromIdentifiers(kfinance_client=mock_client)
         args = GetCapitalizationFromIdentifiersArgs(
@@ -73,26 +87,22 @@ class TestGetCapitalizationFromCompanyIds:
         WHEN we request most recent market caps for multiple companies
         THEN we only get back the most recent market cap for each company
         """
-        expected_response = {
-            "results": {
-                "C_1": {
-                    "capitalizations": [
-                        {
-                            "date": "2024-04-11",
-                            "market_cap": {"unit": "USD", "value": "132416066761.00"},
-                        }
-                    ]
-                },
-                "C_2": {
-                    "capitalizations": [
-                        {
-                            "date": "2024-04-11",
-                            "market_cap": {"unit": "USD", "value": "132416066761.00"},
-                        }
-                    ]
-                },
-            }
-        }
+
+        capitalization = Capitalizations(
+            market_caps=[
+                DailyCapitalization(
+                    date=date(2024, 4, 11),
+                    market_cap=Money(value=Decimal(132416066761), unit="USD"),
+                    tev=None,
+                    shares_outstanding=None,
+                )
+            ]
+        )
+
+        expected_response = GetCapitalizationFromIdentifiersResp(
+            capitalization=Capitalization.market_cap,
+            results={"C_1": capitalization, "C_2": capitalization},
+        )
 
         company_ids = [1, 2]
         for company_id in company_ids:

kfinance/domains/companies/company_tools.py

@@ -27,7 +27,7 @@ class GetInfoFromIdentifiers(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]) -> GetInfoFromIdentifiersResp:
         """Sample response:
 
         {   "results": {
@@ -65,10 +65,9 @@ class GetInfoFromIdentifiers(KfinanceTool):
         info_responses: dict[str, dict] = process_tasks_in_thread_pool_executor(
             api_client=api_client, tasks=tasks
         )
-        resp_model = GetInfoFromIdentifiersResp(
+        return GetInfoFromIdentifiersResp(
             results=info_responses, errors=list(id_triple_resp.errors.values())
         )
-        return resp_model.model_dump(mode="json")
 
 
 class GetCompanyOtherNamesFromIdentifiersResp(ToolRespWithErrors):
@@ -88,7 +87,7 @@ class GetCompanyOtherNamesFromIdentifiers(KfinanceTool):
     def _run(
         self,
         identifiers: list[str],
-    ) -> dict:
+    ) -> GetCompanyOtherNamesFromIdentifiersResp:
         api_client = self.kfinance_client.kfinance_api_client
         id_triple_resp = api_client.unified_fetch_id_triples(identifiers=identifiers)
         tasks = [
@@ -102,10 +101,9 @@ class GetCompanyOtherNamesFromIdentifiers(KfinanceTool):
         info_responses: dict[str, CompanyOtherNames] = process_tasks_in_thread_pool_executor(
             api_client=api_client, tasks=tasks
         )
-        resp_model = GetCompanyOtherNamesFromIdentifiersResp(
+        return GetCompanyOtherNamesFromIdentifiersResp(
             results=info_responses, errors=list(id_triple_resp.errors.values())
         )
-        return resp_model.model_dump(mode="json")
 
 
 class GetCompanySummaryFromIdentifiersResp(ToolRespWithErrors):
@@ -125,7 +123,7 @@ class GetCompanySummaryFromIdentifiers(KfinanceTool):
     def _run(
         self,
         identifiers: list[str],
-    ) -> dict:
+    ) -> GetCompanySummaryFromIdentifiersResp:
         api_client = self.kfinance_client.kfinance_api_client
         id_triple_resp = api_client.unified_fetch_id_triples(identifiers=identifiers)
 
@@ -147,10 +145,9 @@ class GetCompanySummaryFromIdentifiers(KfinanceTool):
             for identifier, descriptions in company_description_responses.items()
         }
 
-        resp_model = GetCompanySummaryFromIdentifiersResp(
+        return GetCompanySummaryFromIdentifiersResp(
             results=summary_results, errors=list(id_triple_resp.errors.values())
         )
-        return resp_model.model_dump(mode="json")
 
 
 class GetCompanyDescriptionFromIdentifiersResp(ToolRespWithErrors):
@@ -170,7 +167,7 @@ class GetCompanyDescriptionFromIdentifiers(KfinanceTool):
     def _run(
         self,
         identifiers: list[str],
-    ) -> dict:
+    ) -> GetCompanyDescriptionFromIdentifiersResp:
         api_client = self.kfinance_client.kfinance_api_client
         id_triple_resp = api_client.unified_fetch_id_triples(identifiers=identifiers)
 
@@ -192,7 +189,6 @@ class GetCompanyDescriptionFromIdentifiers(KfinanceTool):
             for identifier, descriptions in company_description_responses.items()
         }
 
-        resp_model = GetCompanyDescriptionFromIdentifiersResp(
+        return GetCompanyDescriptionFromIdentifiersResp(
             results=description_results, errors=list(id_triple_resp.errors.values())
         )
-        return resp_model.model_dump(mode="json")

kfinance/domains/companies/tests/test_company_tools.py

@@ -4,9 +4,13 @@ from kfinance.client.kfinance import Client
 from kfinance.conftest import SPGI_COMPANY_ID
 from kfinance.domains.companies.company_tools import (
     GetCompanyDescriptionFromIdentifiers,
+    GetCompanyDescriptionFromIdentifiersResp,
     GetCompanyOtherNamesFromIdentifiers,
+    GetCompanyOtherNamesFromIdentifiersResp,
     GetCompanySummaryFromIdentifiers,
+    GetCompanySummaryFromIdentifiersResp,
     GetInfoFromIdentifiers,
+    GetInfoFromIdentifiersResp,
 )
 from kfinance.integrations.tool_calling.tool_calling_models import ToolArgsWithIdentifiers
 
@@ -20,12 +24,14 @@ class TestGetInfoFromIdentifiers:
         """
 
         info_resp = {"name": "S&P Global Inc.", "status": "Operating"}
-        expected_response = {
-            "results": {"SPGI": info_resp},
-            "errors": [
-                "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
-            ],
-        }
+        expected_response = GetInfoFromIdentifiersResp.model_validate(
+            {
+                "results": {"SPGI": info_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/info/{SPGI_COMPANY_ID}",
             json=info_resp,
@@ -61,7 +67,9 @@ class TestGetCompanyDescriptions:
         tool = GetCompanySummaryFromIdentifiers(kfinance_client=mock_client)
         args = ToolArgsWithIdentifiers(identifiers=["SPGI"])
         response = tool.run(args.model_dump(mode="json"))
-        expected_response = {"results": {"SPGI": self.summary}}
+        expected_response = GetCompanySummaryFromIdentifiersResp.model_validate(
+            {"results": {"SPGI": self.summary}}
+        )
         assert response == expected_response
 
     def test_get_company_description_from_identifier(
@@ -81,7 +89,9 @@ class TestGetCompanyDescriptions:
         tool = GetCompanyDescriptionFromIdentifiers(kfinance_client=mock_client)
         args = ToolArgsWithIdentifiers(identifiers=["SPGI"])
         response = tool.run(args.model_dump(mode="json"))
-        expected_response = {"results": {"SPGI": self.description}}
+        expected_response = GetCompanyDescriptionFromIdentifiersResp.model_validate(
+            {"results": {"SPGI": self.description}}
+        )
         assert response == expected_response
 
 
@@ -121,5 +131,7 @@ class TestGetCompanyOtherNames:
         tool = GetCompanyOtherNamesFromIdentifiers(kfinance_client=mock_client)
         args = ToolArgsWithIdentifiers(identifiers=["SPGI"])
         response = tool.run(args.model_dump(mode="json"))
-        expected_response = {"results": {"SPGI": self.company_other_names_info}}
+        expected_response = GetCompanyOtherNamesFromIdentifiersResp.model_validate(
+            {"results": {"SPGI": self.company_other_names_info}}
+        )
         assert response == expected_response

kfinance/domains/competitors/competitor_tools.py

@@ -27,7 +27,7 @@ class GetCompetitorsFromIdentifiers(KfinanceTool):
         self,
         identifiers: list[str],
         competitor_source: CompetitorSource,
-    ) -> dict:
+    ) -> GetCompetitorsFromIdentifiersResp:
         """Sample response:
 
         {
@@ -56,7 +56,6 @@ class GetCompetitorsFromIdentifiers(KfinanceTool):
         competitor_responses: dict[str, CompetitorResponse] = process_tasks_in_thread_pool_executor(
             api_client=api_client, tasks=tasks
         )
-        resp_model = GetCompetitorsFromIdentifiersResp(
+        return GetCompetitorsFromIdentifiersResp(
             results=competitor_responses, errors=list(id_triple_resp.errors.values())
         )
-        return resp_model.model_dump(mode="json")

kfinance/domains/competitors/tests/test_competitor_tools.py

@@ -6,6 +6,7 @@ from kfinance.domains.competitors.competitor_models import CompetitorSource
 from kfinance.domains.competitors.competitor_tools import (
     GetCompetitorsFromIdentifiers,
     GetCompetitorsFromIdentifiersArgs,
+    GetCompetitorsFromIdentifiersResp,
 )
 
 
@@ -24,25 +25,27 @@ class TestGetCompetitorsFromIdentifiers:
                 {"company_id": 4003514, "company_name": "London Stock Exchange Group plc"},
             ]
         }
-        expected_response = {
-            "results": {
-                "SPGI": {
-                    "competitors": [
-                        {
-                            "company_id": "C_35352",
-                            "company_name": "The Descartes Systems Group Inc.",
-                        },
-                        {
-                            "company_id": "C_4003514",
-                            "company_name": "London Stock Exchange Group plc",
-                        },
-                    ]
-                }
-            },
-            "errors": [
-                "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
-            ],
-        }
+        expected_response = GetCompetitorsFromIdentifiersResp.model_validate(
+            {
+                "results": {
+                    "SPGI": {
+                        "competitors": [
+                            {
+                                "company_id": 35352,
+                                "company_name": "The Descartes Systems Group Inc.",
+                            },
+                            {
+                                "company_id": 4003514,
+                                "company_name": "London Stock Exchange Group plc",
+                            },
+                        ]
+                    }
+                },
+                "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/competitors/{SPGI_COMPANY_ID}/named_by_competitor",

kfinance/domains/cusip_and_isin/cusip_and_isin_tools.py

@@ -23,7 +23,7 @@ class GetCusipFromIdentifiers(KfinanceTool):
     args_schema: Type[BaseModel] = ToolArgsWithIdentifiers
     accepted_permissions: set[Permission] | None = {Permission.IDPermission}
 
-    def _run(self, identifiers: list[str]) -> dict[str, str]:
+    def _run(self, identifiers: list[str]) -> GetCusipOrIsinFromIdentifiersResp:
         """Sample response:
 
         {
@@ -45,14 +45,13 @@ class GetCusipFromIdentifiers(KfinanceTool):
         ]
 
         cusip_responses = process_tasks_in_thread_pool_executor(api_client=api_client, tasks=tasks)
-        resp_model = GetCusipOrIsinFromIdentifiersResp(
+        return GetCusipOrIsinFromIdentifiersResp(
             results={
                 identifier: cusip_resp["cusip"]
                 for identifier, cusip_resp in cusip_responses.items()
             },
             errors=list(id_triple_resp.errors.values()),
         )
-        return resp_model.model_dump(mode="json")
 
 
 class GetIsinFromIdentifiers(KfinanceTool):
@@ -61,7 +60,7 @@ class GetIsinFromIdentifiers(KfinanceTool):
     args_schema: Type[BaseModel] = ToolArgsWithIdentifiers
     accepted_permissions: set[Permission] | None = {Permission.IDPermission}
 
-    def _run(self, identifiers: list[str]) -> dict:
+    def _run(self, identifiers: list[str]) -> GetCusipOrIsinFromIdentifiersResp:
         """Sample response:
 
         {
@@ -83,10 +82,9 @@ class GetIsinFromIdentifiers(KfinanceTool):
         ]
 
         isin_responses = process_tasks_in_thread_pool_executor(api_client=api_client, tasks=tasks)
-        resp_model = GetCusipOrIsinFromIdentifiersResp(
+        return GetCusipOrIsinFromIdentifiersResp(
             results={
                 identifier: isin_resp["isin"] for identifier, isin_resp in isin_responses.items()
             },
             errors=list(id_triple_resp.errors.values()),
         )
-        return resp_model.model_dump(mode="json")

kfinance/domains/cusip_and_isin/tests/test_cusip_and_isin_tools.py

@@ -4,6 +4,7 @@ from kfinance.client.kfinance import Client
 from kfinance.conftest import SPGI_SECURITY_ID
 from kfinance.domains.cusip_and_isin.cusip_and_isin_tools import (
     GetCusipFromIdentifiers,
+    GetCusipOrIsinFromIdentifiersResp,
     GetIsinFromIdentifiers,
 )
 from kfinance.integrations.tool_calling.tool_calling_models import ToolArgsWithIdentifiers
@@ -18,10 +19,12 @@ class TestGetCusipFromIdentifiers:
         """
 
         spgi_cusip = "78409V104"
-        expected_response = {
-            "results": {"SPGI": "78409V104"},
-            "errors": ["private_company is a private company without a security_id."],
-        }
+        expected_response = GetCusipOrIsinFromIdentifiersResp.model_validate(
+            {
+                "results": {"SPGI": "78409V104"},
+                "errors": ["private_company is a private company without a security_id."],
+            }
+        )
         requests_mock.get(
             url=f"https://kfinance.kensho.com/api/v1/cusip/{SPGI_SECURITY_ID}",
             json={"cusip": spgi_cusip},
@@ -43,10 +46,12 @@ class TestGetIsinFromIdentifiers:
 
         spgi_isin = "US78409V1044"
 
-        expected_response = {
-            "results": {"SPGI": "US78409V1044"},
-            "errors": ["private_company is a private company without a security_id."],
-        }
+        expected_response = GetCusipOrIsinFromIdentifiersResp.model_validate(
+            {
+                "results": {"SPGI": "US78409V1044"},
+                "errors": ["private_company is a private company without a security_id."],
+            }
+        )
         requests_mock.get(
             url=f"https://kfinance.kensho.com/api/v1/isin/{SPGI_SECURITY_ID}",
             json={"isin": spgi_isin},

kfinance/domains/earnings/earning_tools.py

@@ -35,7 +35,7 @@ class GetEarningsFromIdentifiers(KfinanceTool):
         Permission.TranscriptsPermission,
     }
 
-    def _run(self, identifiers: list[str]) -> dict:
+    def _run(self, identifiers: list[str]) -> GetEarningsFromIdentifiersResp:
         """Sample response:
 
         {
@@ -52,10 +52,9 @@ class GetEarningsFromIdentifiers(KfinanceTool):
         }
 
         """
-        earnings_responses = get_earnings_from_identifiers(
+        return get_earnings_from_identifiers(
             identifiers=identifiers, kfinance_api_client=self.kfinance_client.kfinance_api_client
         )
-        return earnings_responses.model_dump(mode="json")
 
 
 class GetLatestEarningsFromIdentifiers(KfinanceTool):
@@ -71,7 +70,7 @@ class GetLatestEarningsFromIdentifiers(KfinanceTool):
         Permission.TranscriptsPermission,
     }
 
-    def _run(self, identifiers: list[str]) -> dict:
+    def _run(self, identifiers: list[str]) -> GetNextOrLatestEarningsFromIdentifiersResp:
         """Sample response:
 
         {
@@ -95,7 +94,7 @@ class GetLatestEarningsFromIdentifiers(KfinanceTool):
                 output_model.results[identifier] = most_recent_earnings
             else:
                 output_model.errors.append(f"No latest earnings available for {identifier}.")
-        return output_model.model_dump(mode="json")
+        return output_model
 
 
 class GetNextEarningsFromIdentifiers(KfinanceTool):
@@ -111,7 +110,7 @@ class GetNextEarningsFromIdentifiers(KfinanceTool):
         Permission.TranscriptsPermission,
     }
 
-    def _run(self, identifiers: list[str]) -> dict:
+    def _run(self, identifiers: list[str]) -> GetNextOrLatestEarningsFromIdentifiersResp:
         """Sample response:
 
         {
@@ -135,7 +134,7 @@ class GetNextEarningsFromIdentifiers(KfinanceTool):
                 output_model.results[identifier] = next_earnings
             else:
                 output_model.errors.append(f"No next earnings available for {identifier}.")
-        return output_model.model_dump(mode="json")
+        return output_model
 
 
 def get_earnings_from_identifiers(
@@ -170,12 +169,16 @@ class GetTranscriptFromKeyDevIdArgs(BaseModel):
     key_dev_id: int = Field(description="The key dev ID for the earnings call")
 
 
+class GetTranscriptFromKeyDevIdResp(BaseModel):
+    transcript: str
+
+
 class GetTranscriptFromKeyDevId(KfinanceTool):
     name: str = "get_transcript_from_key_dev_id"
     description: str = "Get the raw transcript text for an earnings call by key dev ID."
     args_schema: Type[BaseModel] = GetTranscriptFromKeyDevIdArgs
     accepted_permissions: set[Permission] | None = {Permission.TranscriptsPermission}
 
-    def _run(self, key_dev_id: int) -> str:
+    def _run(self, key_dev_id: int) -> GetTranscriptFromKeyDevIdResp:
         transcript = self.kfinance_client.transcript(key_dev_id)
-        return transcript.raw
+        return GetTranscriptFromKeyDevIdResp(transcript=transcript.raw)