kensho-kfinance 2.9.0__py3-none-any.whl → 3.0.1__py3-none-any.whl

kfinance/domains/segments/segment_tools.py

@@ -0,0 +1,91 @@
+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.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.integrations.tool_calling.tool_calling_models import (
+    KfinanceTool,
+    ToolArgsWithIdentifiers,
+    ValidQuarter,
+)
+
+
+class GetSegmentsFromIdentifiersArgs(ToolArgsWithIdentifiers):
+    # no description because the description for enum fields comes from the enum docstring.
+    segment_type: SegmentType
+    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")
+
+
+class GetSegmentsFromIdentifiers(KfinanceTool):
+    name: str = "get_segments_from_identifiers"
+    description: str = "Get the templated segments associated with a list of identifiers."
+    args_schema: Type[BaseModel] = GetSegmentsFromIdentifiersArgs
+    accepted_permissions: set[Permission] | None = {Permission.SegmentsPermission}
+
+    def _run(
+        self,
+        identifiers: list[str],
+        segment_type: SegmentType,
+        period_type: PeriodType | None = None,
+        start_year: int | None = None,
+        end_year: int | None = None,
+        start_quarter: Literal[1, 2, 3, 4] | None = None,
+        end_quarter: Literal[1, 2, 3, 4] | None = None,
+    ) -> dict:
+        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
+        )
+
+        tasks = [
+            Task(
+                func=api_client.fetch_segments,
+                kwargs=dict(
+                    company_id=company_id,
+                    segment_type=segment_type,
+                    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, company_id in identifiers_to_company_ids.items()
+        ]
+
+        segments_responses = process_tasks_in_thread_pool_executor(
+            api_client=api_client, tasks=tasks
+        )
+
+        # 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
+        ):
+            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}
+
+        return {
+            str(identifier): segments["segments"]
+            for identifier, segments in segments_responses.items()
+        }

kfinance/domains/segments/tests/test_segment_tools.py

@@ -0,0 +1,80 @@
+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,
+)
+
+
+class TestGetSegmentsFromIdentifier:
+    segments_response = {
+        "segments": {
+            "2020": {
+                "Commodity Insights": {
+                    "CAPEX": -7000000.0,
+                    "D&A": 17000000.0,
+                },
+                "Unallocated Assets Held for Sale": None,
+            },
+            "2021": {
+                "Commodity Insights": {
+                    "CAPEX": -2000000.0,
+                    "D&A": 12000000.0,
+                },
+                "Unallocated Assets Held for Sale": {"Total Assets": 321000000.0},
+            },
+        },
+    }
+
+    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
+        """
+
+        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,
+        )
+
+        expected_response = {"SPGI": self.segments_response["segments"]}
+
+        tool = GetSegmentsFromIdentifiers(kfinance_client=mock_client)
+        args = GetSegmentsFromIdentifiersArgs(
+            identifiers=["SPGI"], segment_type=SegmentType.business
+        )
+        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
+        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": {"2021": self.segments_response["segments"]["2021"]},
+            "C_2": {"2021": self.segments_response["segments"]["2021"]},
+        }
+
+        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,
+            )
+
+        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

kfinance/domains/statements/statement_tools.py

@@ -0,0 +1,116 @@
+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.integrations.tool_calling.tool_calling_models import (
+    KfinanceTool,
+    ToolArgsWithIdentifiers,
+)
+
+
+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: Literal[1, 2, 3, 4] | None = Field(default=None, description="Starting quarter")
+    end_quarter: Literal[1, 2, 3, 4] | None = Field(default=None, description="Ending quarter")
+
+
+class GetFinancialStatementFromIdentifiers(KfinanceTool):
+    name: str = "get_financial_statement_from_identifiers"
+    description: str = dedent("""
+        Get a financial statement associated with 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.
+
+        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)
+    """).strip()
+    args_schema: Type[BaseModel] = GetFinancialStatementFromIdentifiersArgs
+    accepted_permissions: set[Permission] | None = {
+        Permission.StatementsPermission,
+        Permission.PrivateCompanyFinancialsPermission,
+    }
+
+    def _run(
+        self,
+        identifiers: list[str],
+        statement: StatementType,
+        period_type: PeriodType | None = None,
+        start_year: int | None = None,
+        end_year: int | None = None,
+        start_quarter: Literal[1, 2, 3, 4] | None = None,
+        end_quarter: Literal[1, 2, 3, 4] | None = None,
+    ) -> dict:
+        """Sample response:
+
+        {
+            'SPGI': {
+                'Revenues': {'2020': 7442000000.0, '2021': 8243000000.0},
+                'Total Revenues': {'2020': 7442000000.0, '2021': 8243000000.0}
+            }
+        }
+        """
+        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
+        )
+
+        tasks = [
+            Task(
+                func=api_client.fetch_statement,
+                kwargs=dict(
+                    company_id=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, company_id in identifiers_to_company_ids.items()
+        ]
+
+        statement_responses = 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

kfinance/domains/statements/tests/test_statement_tools.py

@@ -0,0 +1,73 @@
+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.statements.statement_models import StatementType
+from kfinance.domains.statements.statement_tools import (
+    GetFinancialStatementFromIdentifiers,
+    GetFinancialStatementFromIdentifiersArgs,
+)
+
+
+class TestGetFinancialStatementFromIdentifiers:
+    statement_resp = {
+        "statements": {
+            "2020": {"Revenues": "7442000000.000000", "Total Revenues": "7442000000.000000"},
+            "2021": {"Revenues": "8243000000.000000", "Total Revenues": "8243000000.000000"},
+        }
+    }
+
+    def test_get_financial_statement_from_identifiers(
+        self, mock_client: Client, requests_mock: Mocker
+    ):
+        """
+        GIVEN the GetFinancialLineItemFromIdentifiers tool
+        WHEN we request the SPGI income statement
+        THEN we get back the SPGI income statement
+        """
+
+        requests_mock.get(
+            url=f"https://kfinance.kensho.com/api/v1/statements/{SPGI_COMPANY_ID}/income_statement/none/none/none/none/none",
+            json=self.statement_resp,
+        )
+        expected_response = {
+            "SPGI": {
+                "Revenues": {"2020": 7442000000.0, "2021": 8243000000.0},
+                "Total Revenues": {"2020": 7442000000.0, "2021": 8243000000.0},
+            }
+        }
+
+        tool = GetFinancialStatementFromIdentifiers(kfinance_client=mock_client)
+        args = GetFinancialStatementFromIdentifiersArgs(
+            identifiers=["SPGI"], 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
+        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}},
+        }
+
+        for company_id in company_ids:
+            requests_mock.get(
+                url=f"https://kfinance.kensho.com/api/v1/statements/{company_id}/income_statement/none/none/none/none/none",
+                json=self.statement_resp,
+            )
+
+        tool = GetFinancialStatementFromIdentifiers(kfinance_client=mock_client)
+        args = GetFinancialStatementFromIdentifiersArgs(
+            identifiers=[f"{COMPANY_ID_PREFIX}{company_id}" for company_id in company_ids],
+            statement=StatementType.income_statement,
+        )
+        response = tool.run(args.model_dump(mode="json"))
+        assert response == expected_response

kfinance/integrations/README.md

@@ -0,0 +1,8 @@
+# Integrations
+
+Integrations holds integrations downstream from the main
+Kfinance `Client`. This currently includes tool calling
+and MCP. Note that only shared functionality for tool 
+calling should be put into the `tool_calling` directory.
+The tools themselves should be defined in the relevant
+`domains` sub directory.

kfinance/{mcp.py → integrations/mcp/mcp.py}

@@ -6,8 +6,8 @@ from fastmcp.tools import FunctionTool
 from fastmcp.utilities.logging import get_logger
 from langchain_core.utils.function_calling import convert_to_openai_tool
 
-from kfinance.kfinance import Client
-from kfinance.tool_calling import KfinanceTool
+from kfinance.client.kfinance import Client
+from kfinance.integrations.tool_calling.tool_calling_models import KfinanceTool
 
 
 logger = get_logger(__name__)

kfinance/{tests → integrations/tests}/test_example_notebook.py

@@ -55,7 +55,7 @@ def test_run_notebook(jupyter_kernel_name: str):
     # - mocks for all calls made by the client while executing the notebook
     startup_cell_code = dedent("""
         from datetime import datetime
-        from kfinance.kfinance import Client
+        from kfinance.client.kfinance import Client
         kfinance_client = Client(refresh_token="foo")
         api_client = kfinance_client.kfinance_api_client
         # Set access token so that the client doesn't try to fetch it.
@@ -161,13 +161,13 @@ def test_run_notebook(jupyter_kernel_name: str):
         )
         # Mock out image_open so that we don't have to return an actual png.
         from unittest.mock import MagicMock
-        import kfinance.kfinance
-        kfinance.kfinance.image_open = MagicMock()
+        import kfinance.client.kfinance
+        kfinance.client.kfinance.image_open = MagicMock()
     """)
 
     # Load the notebook
     notebook_path = Path(
-        Path(__file__).parent.parent.parent, "example_notebooks", "basic_usage.ipynb"
+        Path(__file__).parent.parent.parent.parent, "example_notebooks", "basic_usage.ipynb"
     )
     with notebook_path.open() as f:
         nb = nbformat.read(f, as_version=4)

kfinance/{tool_calling → integrations/tool_calling}/README.md

@@ -2,7 +2,7 @@
 
 The tools defined in this directory are intended for tool calling. 
 
-Each tool is a subclass of [KfinanceTool](shared_models.py), which is
+Each tool is a subclass of [KfinanceTool](tool_calling_models.py), which is
 turn a subclass of `BaseTool` from langchain. We use langchain to convert these tools 
 into LLM-specific tool descriptions.
 
@@ -14,7 +14,7 @@ of the arguments.
 - args_schema: A pydantic model defining the input schema for the function (more on that below)
 - _run: the source code for the tool.
 
-All new tools have to be added to the `ALL_TOOLS` [list](__init__.py).
+All new tools have to be added to the `ALL_TOOLS` [list](all_tools.py).
 
 
 When initializing a `KfinanceTool`, it's always required to pass in an initialized Kfinance 

kfinance/integrations/tool_calling/all_tools.py

@@ -0,0 +1,55 @@
+from kfinance.domains.business_relationships.business_relationship_tools import (
+    GetBusinessRelationshipFromIdentifiers,
+)
+from kfinance.domains.capitalizations.capitalization_tools import GetCapitalizationFromIdentifiers
+from kfinance.domains.companies.company_tools import GetInfoFromIdentifiers
+from kfinance.domains.competitors.competitor_tools import GetCompetitorsFromIdentifiers
+from kfinance.domains.cusip_and_isin.cusip_and_isin_tools import (
+    GetCusipFromIdentifiers,
+    GetIsinFromIdentifiers,
+)
+from kfinance.domains.earnings.earning_tools import (
+    GetEarningsFromIdentifiers,
+    GetLatestEarningsFromIdentifiers,
+    GetNextEarningsFromIdentifiers,
+    GetTranscriptFromKeyDevId,
+)
+from kfinance.domains.line_items.line_item_tools import GetFinancialLineItemFromIdentifiers
+from kfinance.domains.prices.price_tools import GetPricesFromIdentifiers
+from kfinance.domains.segments.segment_tools import GetSegmentsFromIdentifiers
+from kfinance.domains.statements.statement_tools import GetFinancialStatementFromIdentifiers
+from kfinance.integrations.tool_calling.static_tools.get_latest import GetLatest
+from kfinance.integrations.tool_calling.static_tools.get_n_quarters_ago import GetNQuartersAgo
+from kfinance.integrations.tool_calling.tool_calling_models import KfinanceTool
+
+
+# A list of all available tools
+ALL_TOOLS: list[type[KfinanceTool]] = [
+    # Static / no API call tools
+    GetLatest,
+    GetNQuartersAgo,
+    # Business Relationships
+    GetBusinessRelationshipFromIdentifiers,
+    # Capitalizations
+    GetCapitalizationFromIdentifiers,
+    # Companies
+    GetInfoFromIdentifiers,
+    # Competitors
+    GetCompetitorsFromIdentifiers,
+    # Cusip and Isin
+    GetCusipFromIdentifiers,
+    GetIsinFromIdentifiers,
+    # Earnings
+    GetEarningsFromIdentifiers,
+    GetLatestEarningsFromIdentifiers,
+    GetNextEarningsFromIdentifiers,
+    GetTranscriptFromKeyDevId,
+    # Line Items
+    GetFinancialLineItemFromIdentifiers,
+    # Prices
+    GetPricesFromIdentifiers,
+    # Segments
+    GetSegmentsFromIdentifiers,
+    # Statements
+    GetFinancialStatementFromIdentifiers,
+]

kfinance/{tool_calling → integrations/tool_calling}/prompts.py

@@ -1,4 +1,5 @@
-from kfinance.tool_calling import GetFinancialLineItemFromIdentifier, GetLatest
+from kfinance.domains.line_items.line_item_tools import GetFinancialLineItemFromIdentifiers
+from kfinance.integrations.tool_calling.static_tools.get_latest import GetLatest
 
 
 BASE_PROMPT = f"""
@@ -7,7 +8,7 @@ BASE_PROMPT = f"""
 
     - Always use the `{GetLatest.model_fields["name"].default}` function when asked about the last or most recent quarter or
     when the time is unspecified in the question.
-    - Try to use `{GetFinancialLineItemFromIdentifier.model_fields["name"].default}` for questions about a company's
+    - Try to use `{GetFinancialLineItemFromIdentifiers.model_fields["name"].default}` for questions about a company's
     finances.
     - If the tools do not respond with data that answers the question, then respond by saying that
     you don't have the data available.

kfinance/integrations/tool_calling/static_tools/README.md

@@ -0,0 +1,4 @@
+# Static Tools
+
+The static_tools directory contains tools that don't have to make 
+any API calls.

kfinance/{tool_calling → integrations/tool_calling/static_tools}/get_latest.py

@@ -2,9 +2,9 @@ from typing import Type
 
 from pydantic import BaseModel, Field
 
-from kfinance.models.date_and_period_models import LatestPeriods
-from kfinance.models.permission_models import Permission
-from kfinance.tool_calling.shared_models import KfinanceTool
+from kfinance.client.models.date_and_period_models import LatestPeriods
+from kfinance.client.permission_models import Permission
+from kfinance.integrations.tool_calling.tool_calling_models import KfinanceTool
 
 
 class GetLatestArgs(BaseModel):

kfinance/{tool_calling → integrations/tool_calling/static_tools}/get_n_quarters_ago.py

@@ -2,9 +2,9 @@ from typing import Type
 
 from pydantic import BaseModel, Field
 
-from kfinance.models.date_and_period_models import YearAndQuarter
-from kfinance.models.permission_models import Permission
-from kfinance.tool_calling.shared_models import KfinanceTool
+from kfinance.client.models.date_and_period_models import YearAndQuarter
+from kfinance.client.permission_models import Permission
+from kfinance.integrations.tool_calling.tool_calling_models import KfinanceTool
 
 
 class GetNQuartersAgoArgs(BaseModel):

kfinance/integrations/tool_calling/static_tools/tests/test_get_lastest.py

@@ -0,0 +1,30 @@
+from datetime import datetime
+
+import time_machine
+
+from kfinance.client.kfinance import Client
+from kfinance.integrations.tool_calling.static_tools.get_latest import GetLatest, GetLatestArgs
+
+
+class TestGetLatest:
+    @time_machine.travel(datetime(2025, 1, 1, 12, tzinfo=datetime.now().astimezone().tzinfo))
+    def test_get_latest(self, mock_client: Client):
+        """
+        GIVEN the GetLatest tool
+        WHEN request latest info
+        THEN we get back latest info
+        """
+
+        expected_resp = {
+            "annual": {"latest_year": 2024},
+            "now": {
+                "current_date": "2025-01-01",
+                "current_month": 1,
+                "current_quarter": 1,
+                "current_year": 2025,
+            },
+            "quarterly": {"latest_quarter": 4, "latest_year": 2024},
+        }
+        tool = GetLatest(kfinance_client=mock_client)
+        resp = tool.run(GetLatestArgs().model_dump(mode="json"))
+        assert resp == expected_resp

kfinance/integrations/tool_calling/static_tools/tests/test_get_n_quarters_ago.py

@@ -0,0 +1,24 @@
+from datetime import datetime
+
+import time_machine
+
+from kfinance.client.kfinance import Client
+from kfinance.integrations.tool_calling.static_tools.get_n_quarters_ago import (
+    GetNQuartersAgo,
+    GetNQuartersAgoArgs,
+)
+
+
+class TestGetNQuartersAgo:
+    @time_machine.travel(datetime(2025, 1, 1, 12, tzinfo=datetime.now().astimezone().tzinfo))
+    def test_get_n_quarters_ago(self, mock_client: Client):
+        """
+        GIVEN the GetNQuartersAgo tool
+        WHEN we request 3 quarters ago
+        THEN we get back 3 quarters ago
+        """
+
+        expected_resp = {"quarter": 2, "year": 2024}
+        tool = GetNQuartersAgo(kfinance_client=mock_client)
+        resp = tool.run(GetNQuartersAgoArgs(n=3).model_dump(mode="json"))
+        assert resp == expected_resp