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

kfinance/domains/segments/tests/test_segment_tools.py

@@ -1,31 +1,50 @@
 from requests_mock import Mocker
 
 from kfinance.client.kfinance import Client
-from kfinance.conftest import SPGI_COMPANY_ID
 from kfinance.domains.companies.company_models import COMPANY_ID_PREFIX
 from kfinance.domains.segments.segment_models import SegmentType
 from kfinance.domains.segments.segment_tools import (
     GetSegmentsFromIdentifiers,
     GetSegmentsFromIdentifiersArgs,
+    GetSegmentsFromIdentifiersResp,
 )
 
 
 class TestGetSegmentsFromIdentifier:
     segments_response = {
-        "segments": {
-            "2020": {
-                "Commodity Insights": {
-                    "CAPEX": -7000000.0,
-                    "D&A": 17000000.0,
-                },
-                "Unallocated Assets Held for Sale": None,
+        "currency": "USD",
+        "periods": {
+            "CY2020": {
+                "period_end_date": "2020-12-31",
+                "num_months": 12,
+                "segments": [
+                    {
+                        "name": "Commodity Insights",
+                        "line_items": [
+                            {"name": "CAPEX", "value": "-7000000.0", "sources": []},
+                            {"name": "D&A", "value": "17000000.0", "sources": []},
+                        ],
+                    }
+                ],
             },
-            "2021": {
-                "Commodity Insights": {
-                    "CAPEX": -2000000.0,
-                    "D&A": 12000000.0,
-                },
-                "Unallocated Assets Held for Sale": {"Total Assets": 321000000.0},
+            "CY2021": {
+                "period_end_date": "2021-12-31",
+                "num_months": 12,
+                "segments": [
+                    {
+                        "name": "Commodity Insights",
+                        "line_items": [
+                            {"name": "CAPEX", "value": "-2000000.0", "sources": []},
+                            {"name": "D&A", "value": "12000000.0", "sources": []},
+                        ],
+                    },
+                    {
+                        "name": "Unallocated Assets Held for Sale",
+                        "line_items": [
+                            {"name": "Total Assets", "value": "321000000.0", "sources": []},
+                        ],
+                    },
+                ],
             },
         },
     }
@@ -37,22 +56,41 @@ class TestGetSegmentsFromIdentifier:
         THEN we get back the SPGI business segment and an error for the non-existent company.
         """
 
-        requests_mock.get(
-            url=f"https://kfinance.kensho.com/api/v1/segments/{SPGI_COMPANY_ID}/business/none/none/none/none/none",
-            # truncated from the original API response
-            json=self.segments_response,
+        # Mock the unified_fetch_id_triples response
+        requests_mock.post(
+            url="https://kfinance.kensho.com/api/v1/ids",
+            json={
+                "identifiers_to_id_triples": {
+                    "SPGI": {
+                        "company_id": 21719,
+                        "security_id": 2629107,
+                        "trading_item_id": 2629108,
+                    }
+                },
+                "errors": {
+                    "NON-EXISTENT": "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
+                },
+            },
+        )
+
+        # Mock the fetch_segments response
+        requests_mock.post(
+            url="https://kfinance.kensho.com/api/v1/segments/",
+            json={"results": {"21719": self.segments_response}, "errors": {}},
         )
 
-        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(
-            identifiers=["SPGI", "non-existent"], segment_type=SegmentType.business
+            identifiers=["SPGI", "NON-EXISTENT"], segment_type=SegmentType.business
         )
         response = tool.run(args.model_dump(mode="json"))
         assert response == expected_response
@@ -60,23 +98,114 @@ class TestGetSegmentsFromIdentifier:
     def test_most_recent_request(self, requests_mock: Mocker, mock_client: Client) -> None:
         """
         GIVEN the GetFinancialLineItemFromIdentifiers tool
-        WHEN we request most recent statement for multiple companies
-        THEN we only get back the most recent statement for each company
+        WHEN we request most recent segment for multiple companies
+        THEN we only get back the most recent segment for each company
         """
 
         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": {
+                        "currency": "USD",
+                        "periods": {
+                            "CY2021": {
+                                "period_end_date": "2021-12-31",
+                                "num_months": 12,
+                                "segments": self.segments_response["periods"]["CY2021"]["segments"],
+                            }
+                        },
+                    },
+                    "C_2": {
+                        "currency": "USD",
+                        "periods": {
+                            "CY2021": {
+                                "period_end_date": "2021-12-31",
+                                "num_months": 12,
+                                "segments": self.segments_response["periods"]["CY2021"]["segments"],
+                            }
+                        },
+                    },
+                }
             }
-        }
+        )
 
-        for company_id in company_ids:
-            requests_mock.get(
-                url=f"https://kfinance.kensho.com/api/v1/segments/{company_id}/business/none/none/none/none/none",
-                json=self.segments_response,
-            )
+        # Mock the unified_fetch_id_triples response
+        requests_mock.post(
+            url="https://kfinance.kensho.com/api/v1/ids",
+            json={
+                "identifiers_to_id_triples": {
+                    "C_1": {"company_id": 1, "security_id": 101, "trading_item_id": 201},
+                    "C_2": {"company_id": 2, "security_id": 102, "trading_item_id": 202},
+                },
+                "errors": {},
+            },
+        )
+
+        # Mock the fetch_segments response
+        requests_mock.post(
+            url="https://kfinance.kensho.com/api/v1/segments/",
+            json={
+                "results": {"1": self.segments_response, "2": self.segments_response},
+                "errors": {},
+            },
+        )
+
+        tool = GetSegmentsFromIdentifiers(kfinance_client=mock_client)
+        args = GetSegmentsFromIdentifiersArgs(
+            identifiers=[f"{COMPANY_ID_PREFIX}{company_id}" for company_id in company_ids],
+            segment_type=SegmentType.business,
+        )
+        response = tool.run(args.model_dump(mode="json"))
+        assert response == expected_response
+
+    def test_empty_most_recent_request(self, requests_mock: Mocker, mock_client: Client) -> None:
+        """
+        GIVEN the GetFinancialLineItemFromIdentifiers tool
+        WHEN we request most recent segment for multiple companies
+        THEN we only get back the most recent segment for each company
+        UNLESS no segments exist
+        """
+
+        company_ids = [1, 2]
+        expected_response = GetSegmentsFromIdentifiersResp.model_validate(
+            {
+                "results": {
+                    "C_1": {"currency": "USD", "periods": {}},
+                    "C_2": {
+                        "currency": "USD",
+                        "periods": {
+                            "CY2021": {
+                                "period_end_date": "2021-12-31",
+                                "num_months": 12,
+                                "segments": self.segments_response["periods"]["CY2021"]["segments"],
+                            }
+                        },
+                    },
+                }
+            }
+        )
+
+        # Mock the unified_fetch_id_triples response
+        requests_mock.post(
+            url="https://kfinance.kensho.com/api/v1/ids",
+            json={
+                "identifiers_to_id_triples": {
+                    "C_1": {"company_id": 1, "security_id": 101, "trading_item_id": 201},
+                    "C_2": {"company_id": 2, "security_id": 102, "trading_item_id": 202},
+                },
+                "errors": {},
+            },
+        )
+
+        # Mock the fetch_segments response with different data for different companies
+        requests_mock.post(
+            url="https://kfinance.kensho.com/api/v1/segments/",
+            json={
+                "results": {"1": {"currency": "USD", "periods": {}}, "2": self.segments_response},
+                "errors": {},
+            },
+        )
 
         tool = GetSegmentsFromIdentifiers(kfinance_client=mock_client)
         args = GetSegmentsFromIdentifiersArgs(

kfinance/domains/statements/statement_models.py

@@ -1,8 +1,10 @@
-from typing import Any
+from datetime import date
 
 from pydantic import BaseModel
 from strenum import StrEnum
 
+from kfinance.domains.line_items.line_item_models import BasePeriodsResp, LineItem
+
 
 class StatementType(StrEnum):
     """The type of financial statement"""
@@ -12,5 +14,17 @@ class StatementType(StrEnum):
     cashflow = "cashflow"
 
 
-class StatementsResp(BaseModel):
-    statements: dict[str, Any]
+class Statement(BaseModel):
+    name: str
+    line_items: list[LineItem]
+
+
+class StatementPeriodData(BaseModel):
+    period_end_date: date
+    num_months: int
+    statements: list[Statement]
+
+
+class StatementsResp(BasePeriodsResp):
+    currency: str | None
+    periods: dict[str, StatementPeriodData]  # period -> statement and period data

kfinance/domains/statements/statement_tools.py

@@ -3,9 +3,9 @@ from typing import Literal, Type
 
 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.models.date_and_period_models import NumPeriods, NumPeriodsBack, PeriodType
 from kfinance.client.permission_models import Permission
+from kfinance.domains.line_items.line_item_models import CalendarType
 from kfinance.domains.statements.statement_models import StatementsResp, StatementType
 from kfinance.integrations.tool_calling.tool_calling_models import (
     KfinanceTool,
@@ -18,27 +18,64 @@ from kfinance.integrations.tool_calling.tool_calling_models import (
 class GetFinancialStatementFromIdentifiersArgs(ToolArgsWithIdentifiers):
     # no description because the description for enum fields comes from the enum docstring.
     statement: StatementType
-    period_type: PeriodType | None = Field(default=None, description="The period type")
-    start_year: int | None = Field(default=None, description="The starting year for the data range")
-    end_year: int | None = Field(default=None, description="The ending year for the data range")
-    start_quarter: ValidQuarter | None = Field(default=None, description="Starting quarter")
-    end_quarter: ValidQuarter | None = Field(default=None, description="Ending quarter")
+    period_type: PeriodType | None = Field(
+        default=None, description="The period type (annual or quarterly)"
+    )
+    start_year: int | None = Field(
+        default=None,
+        description="The starting year for the data range. Use null for the most recent data.",
+    )
+    end_year: int | None = Field(
+        default=None,
+        description="The ending year for the data range. Use null for the most recent data.",
+    )
+    start_quarter: ValidQuarter | None = Field(
+        default=None, description="Starting quarter (1-4). Only used when period_type is quarterly."
+    )
+    end_quarter: ValidQuarter | None = Field(
+        default=None, description="Ending quarter (1-4). Only used when period_type is quarterly."
+    )
+    calendar_type: CalendarType | None = Field(
+        default=None, description="Fiscal year or calendar year"
+    )
+    num_periods: NumPeriods | None = Field(
+        default=None, description="The number of periods to retrieve data for (1-99)"
+    )
+    num_periods_back: NumPeriodsBack | None = Field(
+        default=None,
+        description="The end period of the data range expressed as number of periods back relative to the present period (0-99)",
+    )
 
 
 class GetFinancialStatementFromIdentifiersResp(ToolRespWithErrors):
-    results: dict[str, StatementsResp]
+    results: dict[str, StatementsResp]  # identifier -> response
 
 
 class GetFinancialStatementFromIdentifiers(KfinanceTool):
     name: str = "get_financial_statement_from_identifiers"
     description: str = dedent("""
-        Get a financial statement associated with a group of identifiers.
+        Get a financial statement (balance_sheet, income_statement, or cashflow) for a group of identifiers.
 
-        - To fetch the most recent value for the statement, leave start_year, start_quarter, end_year, and end_quarter as None.
+        - When possible, pass multiple identifiers in a single call rather than making multiple calls.
+        - To fetch the most recent statement, leave start_year, start_quarter, end_year, end_quarter, num_periods, and num_periods_back as null.
+        - The tool accepts an optional calendar_type argument, which can either be 'calendar' or 'fiscal'. If 'calendar' is chosen, then start_year and end_year will filter on calendar year, and the output returned will be in calendar years. If 'fiscal' is chosen (which is the default), then start_year and end_year will filter on fiscal year, and the output returned will be in fiscal years.
+        - To filter by time, use either absolute (start_year, end_year, start_quarter, end_quarter) for specific dates like "in 2023" or "Q2 2021", OR relative (num_periods, num_periods_back) for phrases like "last 3 quarters" or "past five years"—but not both.
 
-        Example:
-        Query: "Fetch the balance sheets of BAC and GS for 2024"
-        Function: get_financial_statement_from_company_ids(identifiers=["BAC", "GS"], statement=StatementType.balance_sheet, start_year=2024, end_year=2024)
+        Examples:
+        Query: "Fetch the balance sheets of Bank of America and Goldman Sachs for 2024"
+        Function: get_financial_statement_from_identifiers(identifiers=["Bank of America", "Goldman Sachs"], statement="balance_sheet", period_type="annual", start_year=2024, end_year=2024)
+
+        Query: "Get income statements for NEE and DUK"
+        Function: get_financial_statement_from_identifiers(identifiers=["NEE", "DUK"], statement="income_statement")
+
+        Query: "Q2 2023 cashflow for XOM"
+        Function: get_financial_statement_from_identifiers(identifiers=["XOM"], statement="cashflow", period_type="quarterly", start_year=2023, end_year=2023, start_quarter=2, end_quarter=2)
+
+        Query: "What is the balance sheet for The New York Times for the past 7 years except for the most recent 2 years?"
+        Function: get_financial_statement_from_identifiers(statement_type="balance_sheet", num_periods=7, num_periods_back=2, identifiers=["NYT"])
+
+        Query: "What are the annual income statement for the calendar years between 2013 and 2016 for Alibaba and Wayfair?"
+        Function: get_financial_statement_from_identifiers(statement_type="income_statement", period_type="annual", calendar_type="calendar", start_year=2013, end_year=2016, identifiers=["BABA", "W"])
     """).strip()
     args_schema: Type[BaseModel] = GetFinancialStatementFromIdentifiersArgs
     accepted_permissions: set[Permission] | None = {
@@ -55,15 +92,61 @@ 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:
+        calendar_type: CalendarType | None = None,
+        num_periods: int | None = None,
+        num_periods_back: int | None = None,
+    ) -> GetFinancialStatementFromIdentifiersResp:
         """Sample response:
 
         {
             'results': {
                 'SPGI': {
-                    'statements': {
-                        '2020': {'Revenues': '7442000000.000000', 'Total Revenues': '7442000000.000000'},
-                        '2021': {'Revenues': '8243000000.000000', 'Total Revenues': '8243000000.000000'}
+                    'currency': 'USD',
+                    'periods': {
+                        'CY2020': {
+                            'period_end_date': '2020-12-31',
+                            'num_months': 12,
+                            'statements': [
+                                {
+                                    'name': 'Income Statement',
+                                    'line_items': [
+                                        {
+                                            'name': 'Revenues',
+                                            'value': 7442000000.0,
+                                            'sources': [
+                                                {
+                                                    'type': 'doc-viewer statement',
+                                                    'url': 'https://www.capitaliq.spglobal.com/...'
+                                                }
+                                            ]
+                                        },
+                                        {
+                                            'name': 'Total Revenues',
+                                            'value': 7442000000.0
+                                        }
+                                    ]
+                                }
+                            ]
+                        },
+                        'CY2021': {
+                            'period_end_date': '2021-12-31',
+                            'num_months': 12,
+                            'statements': [
+                                {
+                                    'name': 'Income Statement',
+                                    'line_items': [
+                                        {
+                                            'name': 'Revenues',
+                                            'value': 8243000000.0
+                                        },
+                                        {
+                                            'name': 'Total Revenues',
+                                            'value': 8243000000.0
+                                        }
+                                    ]
+                                }
+                            ]
+                        }
                     }
                 }
             },
@@ -71,29 +154,29 @@ class GetFinancialStatementFromIdentifiers(KfinanceTool):
         }
         """
         api_client = self.kfinance_client.kfinance_api_client
-        id_triple_resp = api_client.unified_fetch_id_triples(identifiers=identifiers)
-
-        tasks = [
-            Task(
-                func=api_client.fetch_statement,
-                kwargs=dict(
-                    company_id=id_triple.company_id,
-                    statement_type=statement.value,
-                    period_type=period_type,
-                    start_year=start_year,
-                    end_year=end_year,
-                    start_quarter=start_quarter,
-                    end_quarter=end_quarter,
-                ),
-                result_key=identifier,
-            )
-            for identifier, id_triple in id_triple_resp.identifiers_to_id_triples.items()
-        ]
-
-        statement_responses: dict[str, StatementsResp] = process_tasks_in_thread_pool_executor(
-            api_client=api_client, tasks=tasks
+
+        # First resolve identifiers to company IDs
+        ids_response = api_client.unified_fetch_id_triples(identifiers)
+
+        # Call the simplified fetch_statement API with company IDs
+        response = api_client.fetch_statement(
+            company_ids=ids_response.company_ids,
+            statement_type=statement.value,
+            period_type=period_type,
+            start_year=start_year,
+            end_year=end_year,
+            start_quarter=start_quarter,
+            end_quarter=end_quarter,
+            calendar_type=calendar_type,
+            num_periods=num_periods,
+            num_periods_back=num_periods_back,
         )
 
+        identifier_to_results = {}
+        for company_id_str, statement_resp in response.results.items():
+            original_identifier = ids_response.get_identifier_from_company_id(int(company_id_str))
+            identifier_to_results[original_identifier] = statement_resp
+
         # 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.
@@ -102,14 +185,15 @@ class GetFinancialStatementFromIdentifiers(KfinanceTool):
             and end_year is None
             and start_quarter is None
             and end_quarter is None
-            and len(statement_responses) > 1
+            and num_periods is None
+            and num_periods_back is None
+            and len(identifier_to_results) > 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}
+            for result in identifier_to_results.values():
+                result.remove_all_periods_other_than_the_most_recent_one()
+
+        all_errors = list(ids_response.errors.values()) + list(response.errors.values())
 
-        output_model = GetFinancialStatementFromIdentifiersResp(
-            results=statement_responses, errors=list(id_triple_resp.errors.values())
+        return GetFinancialStatementFromIdentifiersResp(
+            results=identifier_to_results, errors=all_errors
         )
-        return output_model.model_dump(mode="json")