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

kfinance/domains/prices/tests/test_price_tools.py

@@ -17,8 +17,8 @@ class TestGetHistoryMetadataFromIdentifiers:
     ):
         """
         GIVEN the GetHistoryMetadataFromIdentifiers tool
-        WHEN we request the history metadata for SPGI
-        THEN we get back SPGI's history metadata
+        WHEN we request the history metadata for SPGI and a non-existent company
+        THEN we get back SPGI's history metadata and an error for the non-existent company.
         """
         metadata_resp = {
             "currency": "USD",
@@ -27,7 +27,12 @@ class TestGetHistoryMetadataFromIdentifiers:
             "instrument_type": "Equity",
             "symbol": "SPGI",
         }
-        expected_resp = {"SPGI": metadata_resp}
+        expected_resp = {
+            "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",
@@ -35,11 +40,13 @@ class TestGetHistoryMetadataFromIdentifiers:
         )
 
         tool = GetHistoryMetadataFromIdentifiers(kfinance_client=mock_client)
-        resp = tool.run(ToolArgsWithIdentifiers(identifiers=["SPGI"]).model_dump(mode="json"))
+        resp = tool.run(
+            ToolArgsWithIdentifiers(identifiers=["SPGI", "non-existent"]).model_dump(mode="json")
+        )
         assert resp == expected_resp
 
 
-class TestPricesFromIdentifiers:
+class TestGetPricesFromIdentifiers:
     prices_resp = {
         "currency": "USD",
         "prices": [
@@ -65,8 +72,8 @@ class TestPricesFromIdentifiers:
     def test_get_prices_from_identifiers(self, mock_client: Client, requests_mock: Mocker):
         """
         GIVEN the GetPricesFromIdentifiers tool
-        WHEN we request prices for SPGI
-        THEN we get back prices for SPGI
+        WHEN we request prices for SPGI and a non-existent company
+        THEN we get back prices for SPGI and an erro for the non-existent company
         """
 
         requests_mock.get(
@@ -74,31 +81,38 @@ class TestPricesFromIdentifiers:
             json=self.prices_resp,
         )
         expected_response = {
-            "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"},
-                    },
-                ]
-            }
+            "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(
-            GetPricesFromIdentifiersArgs(identifiers=["SPGI"]).model_dump(mode="json")
+            GetPricesFromIdentifiersArgs(identifiers=["SPGI", "non-existent"]).model_dump(
+                mode="json"
+            )
         )
         assert response == expected_response
 
@@ -119,18 +133,20 @@ class TestPricesFromIdentifiers:
         expected_single_company_response = {
             "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"},
                 }
             ]
         }
         expected_response = {
-            "C_1": expected_single_company_response,
-            "C_2": expected_single_company_response,
+            "results": {
+                "C_1": expected_single_company_response,
+                "C_2": expected_single_company_response,
+            },
         }
         tool = GetPricesFromIdentifiers(kfinance_client=mock_client)
         response = tool.run(

kfinance/domains/segments/segment_models.py

@@ -1,3 +1,6 @@
+from typing import Any
+
+from pydantic import BaseModel
 from strenum import StrEnum
 
 
@@ -6,3 +9,7 @@ class SegmentType(StrEnum):
 
     business = "business"
     geographic = "geographic"
+
+
+class SegmentsResp(BaseModel):
+    segments: dict[str, Any]

kfinance/domains/segments/segment_tools.py

@@ -5,14 +5,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 PeriodType
 from kfinance.client.permission_models import Permission
-from kfinance.domains.companies.company_identifiers import (
-    fetch_company_ids_from_identifiers,
-    parse_identifiers,
-)
-from kfinance.domains.segments.segment_models import SegmentType
+from kfinance.domains.segments.segment_models import SegmentsResp, SegmentType
 from kfinance.integrations.tool_calling.tool_calling_models import (
     KfinanceTool,
     ToolArgsWithIdentifiers,
+    ToolRespWithErrors,
     ValidQuarter,
 )
 
@@ -27,6 +24,10 @@ class GetSegmentsFromIdentifiersArgs(ToolArgsWithIdentifiers):
     end_quarter: ValidQuarter | None = Field(default=None, description="Ending quarter")
 
 
+class GetSegmentsFromIdentifiersResp(ToolRespWithErrors):
+    results: dict[str, SegmentsResp]
+
+
 class GetSegmentsFromIdentifiers(KfinanceTool):
     name: str = "get_segments_from_identifiers"
     description: str = "Get the templated segments associated with a list of identifiers."
@@ -43,17 +44,33 @@ class GetSegmentsFromIdentifiers(KfinanceTool):
         start_quarter: Literal[1, 2, 3, 4] | None = None,
         end_quarter: Literal[1, 2, 3, 4] | None = None,
     ) -> dict:
+        """Sample Response:
+
+        {
+            'results': {
+                'SPGI': {
+                    'segments': {
+                        '2021': {
+                            'Commodity Insights': {'CAPEX': -2000000.0, 'D&A': 12000000.0},
+                            'Unallocated Assets Held for Sale': {'Total Assets': 321000000.0}
+                        }
+                    }
+                }
+            },
+            '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_segments,
                 kwargs=dict(
-                    company_id=company_id,
+                    company_id=id_triple.company_id,
                     segment_type=segment_type,
                     period_type=period_type,
                     start_year=start_year,
@@ -63,10 +80,10 @@ class GetSegmentsFromIdentifiers(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()
         ]
 
-        segments_responses = process_tasks_in_thread_pool_executor(
+        segments_responses: dict[str, SegmentsResp] = process_tasks_in_thread_pool_executor(
             api_client=api_client, tasks=tasks
         )
 
@@ -78,14 +95,14 @@ class GetSegmentsFromIdentifiers(KfinanceTool):
             and end_year is None
             and start_quarter is None
             and end_quarter is None
-            and len(identifiers) > 1
+            and len(segments_responses) > 1
         ):
             for segments_response in segments_responses.values():
-                most_recent_year = max(segments_response["segments"].keys())
-                most_recent_year_data = segments_response["segments"][most_recent_year]
-                segments_response["segments"] = {most_recent_year: most_recent_year_data}
+                most_recent_year = max(segments_response.segments.keys())
+                most_recent_year_data = segments_response.segments[most_recent_year]
+                segments_response.segments = {most_recent_year: most_recent_year_data}
 
-        return {
-            str(identifier): segments["segments"]
-            for identifier, segments in segments_responses.items()
-        }
+        output_model = 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

@@ -33,8 +33,8 @@ class TestGetSegmentsFromIdentifier:
     def test_get_segments_from_identifier(self, mock_client: Client, requests_mock: Mocker):
         """
         GIVEN the GetSegmentsFromIdentifier tool
-        WHEN we request the SPGI business segment
-        THEN we get back the SPGI business segment
+        WHEN we request the business segment for SPGI and an non-existent company
+        THEN we get back the SPGI business segment and an error for the non-existent company.
         """
 
         requests_mock.get(
@@ -43,11 +43,16 @@ class TestGetSegmentsFromIdentifier:
             json=self.segments_response,
         )
 
-        expected_response = {"SPGI": self.segments_response["segments"]}
+        expected_response = {
+            "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(
-            identifiers=["SPGI"], segment_type=SegmentType.business
+            identifiers=["SPGI", "non-existent"], segment_type=SegmentType.business
         )
         response = tool.run(args.model_dump(mode="json"))
         assert response == expected_response
@@ -61,8 +66,10 @@ class TestGetSegmentsFromIdentifier:
 
         company_ids = [1, 2]
         expected_response = {
-            "C_1": {"2021": self.segments_response["segments"]["2021"]},
-            "C_2": {"2021": self.segments_response["segments"]["2021"]},
+            "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:

kfinance/domains/statements/statement_models.py

@@ -1,3 +1,6 @@
+from typing import Any
+
+from pydantic import BaseModel
 from strenum import StrEnum
 
 
@@ -7,3 +10,7 @@ class StatementType(StrEnum):
     balance_sheet = "balance_sheet"
     income_statement = "income_statement"
     cashflow = "cashflow"
+
+
+class StatementsResp(BaseModel):
+    statements: dict[str, Any]

kfinance/domains/statements/statement_tools.py

@@ -1,21 +1,16 @@
 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 (
-    fetch_company_ids_from_identifiers,
-    parse_identifiers,
-)
-from kfinance.domains.statements.statement_models import StatementType
+from kfinance.domains.statements.statement_models import StatementsResp, StatementType
 from kfinance.integrations.tool_calling.tool_calling_models import (
     KfinanceTool,
     ToolArgsWithIdentifiers,
+    ToolRespWithErrors,
 )
 
 
@@ -29,6 +24,10 @@ class GetFinancialStatementFromIdentifiersArgs(ToolArgsWithIdentifiers):
     end_quarter: Literal[1, 2, 3, 4] | None = Field(default=None, description="Ending quarter")
 
 
+class GetFinancialStatementFromIdentifiersResp(ToolRespWithErrors):
+    results: dict[str, StatementsResp]
+
+
 class GetFinancialStatementFromIdentifiers(KfinanceTool):
     name: str = "get_financial_statement_from_identifiers"
     description: str = dedent("""
@@ -59,23 +58,25 @@ class GetFinancialStatementFromIdentifiers(KfinanceTool):
         """Sample response:
 
         {
-            'SPGI': {
-                'Revenues': {'2020': 7442000000.0, '2021': 8243000000.0},
-                'Total Revenues': {'2020': 7442000000.0, '2021': 8243000000.0}
-            }
+            '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']
         }
         """
         api_client = self.kfinance_client.kfinance_api_client
-        parsed_identifiers = parse_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_statement,
                 kwargs=dict(
-                    company_id=company_id,
+                    company_id=id_triple.company_id,
                     statement_type=statement.value,
                     period_type=period_type,
                     start_year=start_year,
@@ -85,32 +86,29 @@ class GetFinancialStatementFromIdentifiers(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()
         ]
 
-        statement_responses = process_tasks_in_thread_pool_executor(
+        statement_responses: dict[str, StatementsResp] = process_tasks_in_thread_pool_executor(
             api_client=api_client, tasks=tasks
         )
 
-        output = dict()
-        for identifier, result in statement_responses.items():
-            df = (
-                pd.DataFrame(result["statements"])
-                .apply(pd.to_numeric)
-                .replace(np.nan, None)
-                .transpose()
-            )
-            # 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.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(statement_responses) > 1
+        ):
+            for statement_response in statement_responses.values():
+                most_recent_year = max(statement_response.statements.keys())
+                most_recent_year_data = statement_response.statements[most_recent_year]
+                statement_response.statements = {most_recent_year: most_recent_year_data}
+
+        output_model = 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

@@ -23,8 +23,8 @@ class TestGetFinancialStatementFromIdentifiers:
     ):
         """
         GIVEN the GetFinancialLineItemFromIdentifiers tool
-        WHEN we request the SPGI income statement
-        THEN we get back the SPGI income statement
+        WHEN we request the income statement for SPGI and a non-existent company
+        THEN we get back the SPGI income statement and an error for the non-existent company.
         """
 
         requests_mock.get(
@@ -32,30 +32,59 @@ class TestGetFinancialStatementFromIdentifiers:
             json=self.statement_resp,
         )
         expected_response = {
-            "SPGI": {
-                "Revenues": {"2020": 7442000000.0, "2021": 8243000000.0},
-                "Total Revenues": {"2020": 7442000000.0, "2021": 8243000000.0},
-            }
+            "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"
+            ],
         }
 
         tool = GetFinancialStatementFromIdentifiers(kfinance_client=mock_client)
         args = GetFinancialStatementFromIdentifiersArgs(
-            identifiers=["SPGI"], statement=StatementType.income_statement
+            identifiers=["SPGI", "non-existent"], statement=StatementType.income_statement
         )
         response = tool.run(args.model_dump(mode="json"))
         assert response == expected_response
 
     def test_most_recent_request(self, requests_mock: Mocker, mock_client: Client) -> None:
         """
-        GIVEN the GetFinancialLineItemFromIdentifiers tool
+        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
         """
 
         company_ids = [1, 2]
         expected_response = {
-            "C_1": {"Revenues": {"2021": 8243000000.0}, "Total Revenues": {"2021": 8243000000.0}},
-            "C_2": {"Revenues": {"2021": 8243000000.0}, "Total Revenues": {"2021": 8243000000.0}},
+            "results": {
+                "C_1": {
+                    "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:

kfinance/integrations/tool_calling/tests/test_tool_calling_models.py

@@ -24,10 +24,10 @@ class TestGetEndpointsFromToolCallsWithGrounding:
         # truncated from the original
         resp_data = {"name": "S&P Global Inc.", "status": "Operating"}
         resp_endpoint = [
-            "https://kfinance.kensho.com/api/v1/id/SPGI",
+            "https://kfinance.kensho.com/api/v1/ids",
             "https://kfinance.kensho.com/api/v1/info/21719",
         ]
-        expected_resp = {"data": {"SPGI": resp_data}, "endpoint_urls": resp_endpoint}
+        expected_resp = {"data": {"results": {"SPGI": resp_data}}, "endpoint_urls": resp_endpoint}
 
         requests_mock.get(
             url=f"https://kfinance.kensho.com/api/v1/info/{SPGI_COMPANY_ID}",

kfinance/integrations/tool_calling/tool_calling_models.py

@@ -1,7 +1,7 @@
-from typing import Annotated, Any, Literal, Type
+from typing import Annotated, Any, Callable, Dict, Literal, Type
 
 from langchain_core.tools import BaseTool
-from pydantic import BaseModel, BeforeValidator, ConfigDict, Field
+from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, model_serializer
 
 from kfinance.client.kfinance import Client
 from kfinance.client.permission_models import Permission
@@ -65,9 +65,7 @@ class KfinanceTool(BaseTool):
         Where feasible and useful, tools should use batch processing to parallelize
         requests, usually by allowing callers to pass in multiple identifiers.
         The usual processing order for functions with multiple identifiers is:
-        - Parse string identifiers into `CompanyIdentifier`
-        - batch fetch the required id for the api call (company/security/trading item id) via
-            fetch_<relevant>_ids_from_identifiers
+        - batch fetch id triples via unified_fetch_id_triples
         - batch fetch the required info based on the ids via process_tasks_in_thread_pool_executor
         - format results for output
 
@@ -119,3 +117,24 @@ def convert_str_to_int(v: Any) -> Any:
 # ValidationError during deserialization unless they have been converted
 # to int.
 ValidQuarter = Annotated[Literal[1, 2, 3, 4], BeforeValidator(convert_str_to_int)]
+
+
+class ToolRespWithErrors(BaseModel):
+    """A tool response with an `errors` field.
+
+    - `errors` is always the last field in the response.
+    - `errors` is only included if there is at least one error.
+    """
+
+    errors: list[str] = Field(default_factory=list)
+
+    @model_serializer(mode="wrap")
+    def serialize_model(self, handler: Callable) -> Dict[str, Any]:
+        """Make `errors` the last response field and only include if there is at least one error."""
+        data = handler(self)
+        errors = data.pop("errors")
+        # data = copy(data)
+        # data.keys()
+        if errors:
+            data["errors"] = errors
+        return data

kfinance/version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.0.3'
-__version_tuple__ = version_tuple = (3, 0, 3)
+__version__ = version = '3.1.0'
+__version_tuple__ = version_tuple = (3, 1, 0)