kensho-kfinance 3.0.3__py3-none-any.whl → 3.1.0__py3-none-any.whl

kfinance/domains/line_items/line_item_tools.py

@@ -1,18 +1,11 @@
 from textwrap import dedent
 from typing import Literal, Type
 
-import numpy as np
-import pandas as pd
 from pydantic import BaseModel, Field
 
 from kfinance.client.batch_request_handling import Task, process_tasks_in_thread_pool_executor
 from kfinance.client.models.date_and_period_models import PeriodType
 from kfinance.client.permission_models import Permission
-from kfinance.domains.companies.company_identifiers import (
-    Identifier,
-    fetch_company_ids_from_identifiers,
-    parse_identifiers,
-)
 from kfinance.domains.line_items.line_item_models import (
     LINE_ITEM_NAMES_AND_ALIASES,
     LineItemResponse,
@@ -20,6 +13,7 @@ from kfinance.domains.line_items.line_item_models import (
 from kfinance.integrations.tool_calling.tool_calling_models import (
     KfinanceTool,
     ToolArgsWithIdentifiers,
+    ToolRespWithErrors,
 )
 
 
@@ -37,6 +31,10 @@ class GetFinancialLineItemFromIdentifiersArgs(ToolArgsWithIdentifiers):
     end_quarter: Literal[1, 2, 3, 4] | None = Field(default=None, description="Ending quarter")
 
 
+class GetFinancialLineItemFromIdentifiersResp(ToolRespWithErrors):
+    results: dict[str, LineItemResponse]
+
+
 class GetFinancialLineItemFromIdentifiers(KfinanceTool):
     name: str = "get_financial_line_item_from_identifiers"
     description: str = dedent("""
@@ -76,16 +74,13 @@ class GetFinancialLineItemFromIdentifiers(KfinanceTool):
         }
         """
         api_client = self.kfinance_client.kfinance_api_client
-        parsed_identifiers = parse_identifiers(identifiers=identifiers, api_client=api_client)
-        identifiers_to_company_ids = fetch_company_ids_from_identifiers(
-            identifiers=parsed_identifiers, api_client=api_client
-        )
+        id_triple_resp = api_client.unified_fetch_id_triples(identifiers=identifiers)
 
         tasks = [
             Task(
                 func=api_client.fetch_line_item,
                 kwargs=dict(
-                    company_id=company_id,
+                    company_id=id_triple.company_id,
                     line_item=line_item,
                     period_type=period_type,
                     start_year=start_year,
@@ -95,31 +90,29 @@ class GetFinancialLineItemFromIdentifiers(KfinanceTool):
                 ),
                 result_key=identifier,
             )
-            for identifier, company_id in identifiers_to_company_ids.items()
+            for identifier, id_triple in id_triple_resp.identifiers_to_id_triples.items()
         ]
 
-        line_item_responses: dict[Identifier, LineItemResponse] = (
-            process_tasks_in_thread_pool_executor(api_client=api_client, tasks=tasks)
+        line_item_responses: dict[str, LineItemResponse] = process_tasks_in_thread_pool_executor(
+            api_client=api_client, tasks=tasks
         )
 
-        output = dict()
-        for identifier, result in line_item_responses.items():
-            df = (
-                pd.DataFrame({"line_item": result.line_item})
-                .apply(pd.to_numeric)
-                .replace(np.nan, None)
-            )
-            # If no date and multiple companies, only return the most recent value.
-            # By default, we return 5 years of data, which can be too much when
-            # returning data for many companies.
-            if (
-                start_year is None
-                and end_year is None
-                and start_quarter is None
-                and end_quarter is None
-                and len(identifiers) > 1
-            ):
-                df = df.tail(1)
-            output[str(identifier)] = df.transpose().set_index(pd.Index([line_item])).to_dict()
-
-        return output
+        # If no date and multiple companies, only return the most recent value.
+        # By default, we return 5 years of data, which can be too much when
+        # returning data for many companies.
+        if (
+            start_year is None
+            and end_year is None
+            and start_quarter is None
+            and end_quarter is None
+            and len(line_item_responses) > 1
+        ):
+            for line_item_response in line_item_responses.values():
+                most_recent_year = max(line_item_response.line_item.keys())
+                most_recent_year_data = line_item_response.line_item[most_recent_year]
+                line_item_response.line_item = {most_recent_year: most_recent_year_data}
+
+        output_model = GetFinancialLineItemFromIdentifiersResp(
+            results=line_item_responses, errors=list(id_triple_resp.errors.values())
+        )
+        return output_model.model_dump(mode="json")

kfinance/domains/line_items/tests/test_line_item_tools.py

@@ -24,8 +24,8 @@ class TestGetFinancialLineItemFromCompanyIds:
     ):
         """
         GIVEN the GetFinancialLineItemFromCompanyId tool
-        WHEN we request SPGI revenue
-        THEN we get back the SPGI revenue
+        WHEN we request revenue for SPGI and a non-existent company
+        THEN we get back the SPGI revenue and an error for the non-existent company
         """
 
         expected_response = {
@@ -35,6 +35,20 @@ class TestGetFinancialLineItemFromCompanyIds:
                 "2024": {"revenue": 14208000000.0},
             }
         }
+        expected_response = {
+            "results": {
+                "SPGI": {
+                    "line_item": {
+                        "2022": "11181000000.000000",
+                        "2023": "12497000000.000000",
+                        "2024": "14208000000.000000",
+                    }
+                }
+            },
+            "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/line_item/{SPGI_COMPANY_ID}/revenue/none/none/none/none/none",
@@ -42,7 +56,9 @@ class TestGetFinancialLineItemFromCompanyIds:
         )
 
         tool = GetFinancialLineItemFromIdentifiers(kfinance_client=mock_client)
-        args = GetFinancialLineItemFromIdentifiersArgs(identifiers=["SPGI"], line_item="revenue")
+        args = GetFinancialLineItemFromIdentifiersArgs(
+            identifiers=["SPGI", "non-existent"], line_item="revenue"
+        )
         response = tool.run(args.model_dump(mode="json"))
         assert response == expected_response
 
@@ -55,8 +71,10 @@ class TestGetFinancialLineItemFromCompanyIds:
 
         company_ids = [1, 2]
         expected_response = {
-            "C_1": {"2024": {"revenue": 14208000000.0}},
-            "C_2": {"2024": {"revenue": 14208000000.0}},
+            "results": {
+                "C_1": {"line_item": {"2024": "14208000000.000000"}},
+                "C_2": {"line_item": {"2024": "14208000000.000000"}},
+            }
         }
         for company_id in company_ids:
             requests_mock.get(

kfinance/domains/mergers_and_acquisitions/merger_and_acquisition_models.py

@@ -0,0 +1,15 @@
+from datetime import date
+
+from pydantic import BaseModel
+
+
+class MergerSummary(BaseModel):
+    transaction_id: int
+    merger_title: str
+    closed_date: date
+
+
+class MergersResp(BaseModel):
+    target: list[MergerSummary]
+    buyer: list[MergerSummary]
+    seller: list[MergerSummary]

kfinance/domains/mergers_and_acquisitions/merger_and_acquisition_tools.py

@@ -6,19 +6,21 @@ from pydantic import BaseModel, Field
 from kfinance.client.batch_request_handling import Task, process_tasks_in_thread_pool_executor
 from kfinance.client.kfinance import Company, MergerOrAcquisition, ParticipantInMerger
 from kfinance.client.permission_models import Permission
-from kfinance.domains.companies.company_identifiers import (
-    CompanyId,
-    fetch_company_ids_from_identifiers,
-    parse_identifiers,
-)
+from kfinance.domains.companies.company_models import prefix_company_id
+from kfinance.domains.mergers_and_acquisitions.merger_and_acquisition_models import MergersResp
 from kfinance.integrations.tool_calling.tool_calling_models import (
     KfinanceTool,
     ToolArgsWithIdentifier,
     ToolArgsWithIdentifiers,
+    ToolRespWithErrors,
 )
 
 
-class GetMergersFromIdentifier(KfinanceTool):
+class GetMergersFromIdentifiersResp(ToolRespWithErrors):
+    results: dict[str, MergersResp]
+
+
+class GetMergersFromIdentifiers(KfinanceTool):
     name: str = "get_mergers_from_identifiers"
     description: str = dedent("""
         Get the transaction IDs that involve the given identifiers.
@@ -29,24 +31,57 @@ class GetMergersFromIdentifier(KfinanceTool):
     accepted_permissions: set[Permission] | None = {Permission.MergersPermission}
 
     def _run(self, identifiers: list[str]) -> dict:
+        """Sample Response:
+
+        {
+            'results': {
+                'SPGI': {
+                    'target': [
+                        {
+                            'transaction_id': 10998717,
+                            'merger_title': 'Closed M/A of Microsoft Corporation',
+                            'closed_date': '2021-01-01'
+                        }
+                    ],
+                    'buyer': [
+                        {
+                           'transaction_id': 517414,
+                           'merger_title': 'Closed M/A of MongoMusic, Inc.',
+                           'closed_date': '2023-01-01'
+                        },
+                    'seller': [
+                        {
+                            'transaction_id': 455551,
+                            'merger_title': 'Closed M/A of VacationSpot.com, Inc.',
+                            'closed_date': '2024-01-01'
+                        },
+                    ]
+                }
+            },
+            'errors': ['No identification triple found for the provided identifier: NON-EXISTENT of type: ticker']
+        }
+
+        """
+
         api_client = self.kfinance_client.kfinance_api_client
-        parsed_identifiers = parse_identifiers(identifiers=identifiers, api_client=api_client)
-        identifiers_to_company_ids = fetch_company_ids_from_identifiers(
-            identifiers=parsed_identifiers, api_client=api_client
-        )
+        id_triple_resp = api_client.unified_fetch_id_triples(identifiers=identifiers)
 
         tasks = [
             Task(
                 func=api_client.fetch_mergers_for_company,
-                kwargs=dict(company_id=company_id),
+                kwargs=dict(company_id=id_triple.company_id),
                 result_key=identifier,
             )
-            for identifier, company_id in identifiers_to_company_ids.items()
+            for identifier, id_triple in id_triple_resp.identifiers_to_id_triples.items()
         ]
 
-        merger_responses = process_tasks_in_thread_pool_executor(api_client=api_client, tasks=tasks)
-
-        return {str(identifier): mergers for identifier, mergers in merger_responses.items()}
+        merger_responses: dict[str, MergersResp] = process_tasks_in_thread_pool_executor(
+            api_client=api_client, tasks=tasks
+        )
+        output_model = GetMergersFromIdentifiersResp(
+            results=merger_responses, errors=list(id_triple_resp.errors.values())
+        )
+        return output_model.model_dump(mode="json")
 
 
 class GetMergerInfoFromTransactionIdArgs(BaseModel):
@@ -83,34 +118,21 @@ class GetMergerInfoFromTransactionId(KfinanceTool):
             else None,
             "participants": {
                 "target": {
-                    "company_id": str(
-                        CompanyId(
-                            company_id=merger_participants["target"].company.company_id,
-                            api_client=self.kfinance_client.kfinance_api_client,
-                        )
+                    "company_id": prefix_company_id(
+                        merger_participants["target"].company.company_id
                     ),
                     "company_name": merger_participants["target"].company.name,
                 },
                 "buyers": [
                     {
-                        "company_id": str(
-                            CompanyId(
-                                buyer.company.company_id,
-                                api_client=self.kfinance_client.kfinance_api_client,
-                            )
-                        ),
+                        "company_id": prefix_company_id(buyer.company.company_id),
                         "company_name": buyer.company.name,
                     }
                     for buyer in merger_participants["buyers"]
                 ],
                 "sellers": [
                     {
-                        "company_id": str(
-                            CompanyId(
-                                seller.company.company_id,
-                                api_client=self.kfinance_client.kfinance_api_client,
-                            )
-                        ),
+                        "company_id": prefix_company_id(seller.company.company_id),
                         "company_name": seller.company.name,
                     }
                     for seller in merger_participants["sellers"]
@@ -161,12 +183,7 @@ class GetAdvisorsForCompanyInTransactionFromIdentifier(KfinanceTool):
         if advisors:
             return [
                 {
-                    "advisor_company_id": str(
-                        CompanyId(
-                            company_id=advisor.company.company_id,
-                            api_client=self.kfinance_client.kfinance_api_client,
-                        )
-                    ),
+                    "advisor_company_id": prefix_company_id(advisor.company.company_id),
                     "advisor_company_name": advisor.company.name,
                     "advisor_type_name": advisor.advisor_type_name,
                 }

kfinance/domains/mergers_and_acquisitions/tests/test_merger_and_acquisition_tools.py

@@ -3,30 +3,43 @@ from copy import deepcopy
 from requests_mock import Mocker
 
 from kfinance.client.kfinance import Client
-from kfinance.client.tests.test_objects import MOCK_COMPANY_DB, ordered
+from kfinance.client.tests.test_objects import (
+    MERGERS_RESP,
+    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_tools import (
     GetAdvisorsForCompanyInTransactionFromIdentifier,
     GetAdvisorsForCompanyInTransactionFromIdentifierArgs,
     GetMergerInfoFromTransactionId,
     GetMergerInfoFromTransactionIdArgs,
-    GetMergersFromIdentifier,
+    GetMergersFromIdentifiers,
 )
 from kfinance.integrations.tool_calling.tool_calling_models import ToolArgsWithIdentifiers
 
 
 class TestGetMergersFromIdentifiers:
     def test_get_mergers_from_identifiers(self, requests_mock: Mocker, mock_client: Client):
-        msft_mergers = MOCK_COMPANY_DB["21835"]["mergers"]
-        expected_response = {"MSFT": msft_mergers}
-        company_id = 21835
+        """
+        GIVEN the GetMergersFromIdentifiers tool
+        WHEN we request mergers for SPGI and a non-existent company
+        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"
+            ],
+        }
         requests_mock.get(
-            url=f"https://kfinance.kensho.com/api/v1/mergers/{company_id}", json=msft_mergers
+            url=f"https://kfinance.kensho.com/api/v1/mergers/{SPGI_COMPANY_ID}", json=merger_data
         )
-        tool = GetMergersFromIdentifier(kfinance_client=mock_client)
-        args = ToolArgsWithIdentifiers(identifiers=["MSFT"])
+        tool = GetMergersFromIdentifiers(kfinance_client=mock_client)
+        args = ToolArgsWithIdentifiers(identifiers=["SPGI", "non-existent"])
         response = tool.run(args.model_dump(mode="json"))
-        assert ordered(response) == ordered(expected_response)
+        assert response == expected_response
 
 
 class TestGetCompaniesAdvisingCompanyInTransactionFromIdentifier:

kfinance/domains/prices/price_models.py

@@ -1,20 +1,12 @@
 from copy import deepcopy
 from datetime import date
-from typing import Any, TypedDict
+from typing import Any
 
 from pydantic import BaseModel, model_validator
 
 from kfinance.client.models.decimal_with_unit import Money, Shares
 
 
-class HistoryMetadata(TypedDict):
-    currency: str
-    symbol: str
-    exchange_name: str
-    instrument_type: str
-    first_trade_date: date
-
-
 class Prices(BaseModel):
     """Prices represents prices for a stock for a specific "date".
 
@@ -68,3 +60,11 @@ class PriceHistory(BaseModel):
                 for key in ["open", "high", "low", "close"]:
                     capitalization[key] = dict(unit=currency, value=capitalization[key])
         return data
+
+
+class HistoryMetadataResp(BaseModel):
+    currency: str
+    symbol: str
+    exchange_name: str
+    instrument_type: str
+    first_trade_date: date

kfinance/domains/prices/price_tools.py

@@ -7,13 +7,11 @@ from pydantic import BaseModel, Field
 from kfinance.client.batch_request_handling import Task, process_tasks_in_thread_pool_executor
 from kfinance.client.models.date_and_period_models import Periodicity
 from kfinance.client.permission_models import Permission
-from kfinance.domains.companies.company_identifiers import (
-    fetch_trading_item_ids_from_identifiers,
-    parse_identifiers,
-)
+from kfinance.domains.prices.price_models import HistoryMetadataResp, PriceHistory
 from kfinance.integrations.tool_calling.tool_calling_models import (
     KfinanceTool,
     ToolArgsWithIdentifiers,
+    ToolRespWithErrors,
 )
 
 
@@ -32,6 +30,10 @@ class GetPricesFromIdentifiersArgs(ToolArgsWithIdentifiers):
     )
 
 
+class GetPricesFromIdentifiersResp(ToolRespWithErrors):
+    results: dict[str, PriceHistory]
+
+
 class GetPricesFromIdentifiers(KfinanceTool):
     name: str = "get_prices_from_identifiers"
     description: str = dedent("""
@@ -81,20 +83,20 @@ class GetPricesFromIdentifiers(KfinanceTool):
                         'volume': {'value': '1182229', 'unit': 'Shares'}
                     }
                 ]
-            }
+            },
+            'errors': ['No identification triple found for the provided identifier: NON-EXISTENT of type: ticker']
         }
         """
+
         api_client = self.kfinance_client.kfinance_api_client
-        parsed_identifiers = parse_identifiers(identifiers=identifiers, api_client=api_client)
-        identifiers_to_trading_item_ids = fetch_trading_item_ids_from_identifiers(
-            identifiers=parsed_identifiers, api_client=api_client
-        )
+        id_triple_resp = api_client.unified_fetch_id_triples(identifiers=identifiers)
+        id_triple_resp.filter_out_companies_without_trading_item_ids()
 
         tasks = [
             Task(
                 func=api_client.fetch_history,
                 kwargs=dict(
-                    trading_item_id=trading_item_id,
+                    trading_item_id=id_triple.trading_item_id,
                     start_date=start_date,
                     end_date=end_date,
                     periodicity=periodicity,
@@ -102,20 +104,26 @@ class GetPricesFromIdentifiers(KfinanceTool):
                 ),
                 result_key=identifier,
             )
-            for identifier, trading_item_id in identifiers_to_trading_item_ids.items()
+            for identifier, id_triple in id_triple_resp.identifiers_to_id_triples.items()
         ]
 
-        price_responses = process_tasks_in_thread_pool_executor(api_client=api_client, tasks=tasks)
+        price_responses: dict[str, PriceHistory] = process_tasks_in_thread_pool_executor(
+            api_client=api_client, tasks=tasks
+        )
+        # If we return results for more than one company and the start and end dates are unset,
+        # truncate data to only return the most recent datapoint.
+        if len(price_responses) > 1 and start_date is None and end_date is None:
+            for price_response in price_responses.values():
+                price_response.prices = price_response.prices[-1:]
+
+        output_model = GetPricesFromIdentifiersResp(
+            results=price_responses, errors=list(id_triple_resp.errors.values())
+        )
+        return output_model.model_dump(mode="json")
 
-        # Only include most recent price if more than one identifier passed and start_date == end_date == None
-        dump_include_filter = None
-        if len(identifiers) > 1 and start_date == end_date is None:
-            dump_include_filter = {"prices": {0: True}}
 
-        return {
-            str(identifier): prices.model_dump(mode="json", include=dump_include_filter)
-            for identifier, prices in price_responses.items()
-        }
+class GetHistoryMetadataFromIdentifiersResp(ToolRespWithErrors):
+    results: dict[str, HistoryMetadataResp]
 
 
 class GetHistoryMetadataFromIdentifiers(KfinanceTool):
@@ -132,34 +140,37 @@ class GetHistoryMetadataFromIdentifiers(KfinanceTool):
         """Sample response:
 
         {
-            'SPGI': {
-                'currency': 'USD',
-                'exchange_name': 'NYSE',
-                'first_trade_date': '1968-01-02',
-                'instrument_type': 'Equity',
-                'symbol': 'SPGI'
-            }
+            'results': {
+                'SPGI': {
+                    'currency': 'USD',
+                    'exchange_name': 'NYSE',
+                    'first_trade_date': '1968-01-02',
+                    'instrument_type': 'Equity',
+                    'symbol': 'SPGI'
+                }
+            },
+            'errors': ['No identification triple found for the provided identifier: NON-EXISTENT of type: ticker']
         }
         """
+
         api_client = self.kfinance_client.kfinance_api_client
-        parsed_identifiers = parse_identifiers(identifiers=identifiers, api_client=api_client)
-        identifiers_to_trading_item_ids = fetch_trading_item_ids_from_identifiers(
-            identifiers=parsed_identifiers, api_client=api_client
-        )
+        id_triple_resp = api_client.unified_fetch_id_triples(identifiers=identifiers)
+        id_triple_resp.filter_out_companies_without_trading_item_ids()
 
         tasks = [
             Task(
                 func=api_client.fetch_history_metadata,
-                kwargs=dict(trading_item_id=trading_item_id),
+                kwargs=dict(trading_item_id=id_triple.trading_item_id),
                 result_key=identifier,
             )
-            for identifier, trading_item_id in identifiers_to_trading_item_ids.items()
+            for identifier, id_triple in id_triple_resp.identifiers_to_id_triples.items()
         ]
 
-        history_metadata_responses = process_tasks_in_thread_pool_executor(
-            api_client=api_client, tasks=tasks
+        history_metadata_responses: dict[str, HistoryMetadataResp] = (
+            process_tasks_in_thread_pool_executor(api_client=api_client, tasks=tasks)
+        )
+        output_model = GetHistoryMetadataFromIdentifiersResp(
+            results=history_metadata_responses, errors=list(id_triple_resp.errors.values())
         )
 
-        return {
-            str(identifier): result for identifier, result in history_metadata_responses.items()
-        }
+        return output_model.model_dump(mode="json")