kensho-kfinance 2.8.0__py3-none-any.whl → 3.0.0__py3-none-any.whl

kfinance/domains/capitalizations/capitalization_tools.py

@@ -0,0 +1,89 @@
+from datetime import date
+from textwrap import dedent
+from typing import Type
+
+from pydantic import BaseModel, Field
+
+from kfinance.client.batch_request_handling import Task, process_tasks_in_thread_pool_executor
+from kfinance.client.permission_models import Permission
+from kfinance.domains.capitalizations.capitalization_models import Capitalization
+from kfinance.domains.companies.company_identifiers import (
+    fetch_company_ids_from_identifiers,
+    parse_identifiers,
+)
+from kfinance.integrations.tool_calling.tool_calling_models import (
+    KfinanceTool,
+    ToolArgsWithIdentifiers,
+)
+
+
+class GetCapitalizationFromIdentifiersArgs(ToolArgsWithIdentifiers):
+    # no description because the description for enum fields comes from the enum docstring.
+    capitalization: Capitalization
+    start_date: date | None = Field(
+        description="The start date for historical capitalization retrieval", default=None
+    )
+    end_date: date | None = Field(
+        description="The end date for historical capitalization retrieval", default=None
+    )
+
+
+class GetCapitalizationFromIdentifiers(KfinanceTool):
+    name: str = "get_capitalization_from_identifiers"
+    description: str = dedent("""
+        Get the historical market cap, tev (Total Enterprise Value), or shares outstanding for a group of identifiers between inclusive start_date and inclusive end date.
+
+        - When possible, pass multiple identifiers in a single call rather than making multiple calls.
+        - When requesting the most recent values, leave start_date and end_date empty.
+
+        Example:
+        Query: "What are the market caps of AAPL and WMT?"
+        Function: get_capitalization_from_identifiers(capitalization=Capitalization.market_cap, identifiers=["AAPL", "WMT"])
+    """).strip()
+    args_schema: Type[BaseModel] = GetCapitalizationFromIdentifiersArgs
+    accepted_permissions: set[Permission] | None = {Permission.PricingPermission}
+
+    def _run(
+        self,
+        identifiers: list[str],
+        capitalization: Capitalization,
+        start_date: str | None = None,
+        end_date: str | None = None,
+    ) -> dict:
+        """Sample response:
+
+        {
+            'SPGI': [
+                {'date': '2024-04-10', 'market_cap': {'unit': 'USD', 'value': '132766738270.00'}},
+                {'date': '2024-04-11', 'market_cap': {'unit': 'USD', 'value': '132416066761.00'}}
+            ]
+        }
+        """
+        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_market_caps_tevs_and_shares_outstanding,
+                kwargs=dict(company_id=company_id, start_date=start_date, end_date=end_date),
+                result_key=identifier,
+            )
+            for identifier, company_id in identifiers_to_company_ids.items()
+        ]
+
+        capitalization_responses = process_tasks_in_thread_pool_executor(
+            api_client=api_client, tasks=tasks
+        )
+
+        return {
+            str(identifier): capitalization_response.model_dump_json_single_metric(
+                capitalization_metric=capitalization,
+                only_include_most_recent_value=True
+                if (len(identifiers) > 1 and start_date == end_date is None)
+                else False,
+            )
+            for identifier, capitalization_response in capitalization_responses.items()
+        }

kfinance/{tests/test_models → domains/capitalizations/tests}/test_capitalization_models.py

@@ -1,8 +1,8 @@
 from datetime import date
 from decimal import Decimal
 
-from kfinance.decimal_with_unit import Money, Shares
-from kfinance.models.capitalization_models import (
+from kfinance.client.models.decimal_with_unit import Money, Shares
+from kfinance.domains.capitalizations.capitalization_models import (
     Capitalization,
     Capitalizations,
     DailyCapitalization,
@@ -69,15 +69,13 @@ class TestCapitalizations:
     def test_single_attribute_serialization(self):
         """
         GIVEN a Capitalizations object
-        WHEN we only want to jsonify a single attribute like market_cap
+        WHEN we only want to dump a single attribute like market_cap
         THEN that attribute gets returned in an LLM-readable, jsonifyable dict format
         """
-        expected_serialization = {
-            "market_cap": [
-                {"2024-06-24": {"unit": "USD", "value": "139231113000.00"}},
-                {"2024-06-25": {"unit": "USD", "value": "140423262000.00"}},
-            ]
-        }
+        expected_serialization = [
+            {"date": "2024-06-24", "market_cap": {"unit": "USD", "value": "139231113000.00"}},
+            {"date": "2024-06-25", "market_cap": {"unit": "USD", "value": "140423262000.00"}},
+        ]
         capitalizations = Capitalizations.model_validate(self.api_resp)
-        serialized = capitalizations.jsonify_single_attribute(Capitalization.market_cap)
+        serialized = capitalizations.model_dump_json_single_metric(Capitalization.market_cap)
         assert serialized == expected_serialization

kfinance/domains/capitalizations/tests/test_capitalization_tools.py

@@ -0,0 +1,85 @@
+from requests_mock import Mocker
+
+from kfinance.client.kfinance import Client
+from kfinance.conftest import SPGI_COMPANY_ID
+from kfinance.domains.capitalizations.capitalization_models import Capitalization
+from kfinance.domains.capitalizations.capitalization_tools import (
+    GetCapitalizationFromIdentifiers,
+    GetCapitalizationFromIdentifiersArgs,
+)
+from kfinance.domains.companies.company_models import COMPANY_ID_PREFIX
+
+
+class TestGetCapitalizationFromCompanyIds:
+    market_caps_resp = {
+        "currency": "USD",
+        "market_caps": [
+            {
+                "date": "2024-04-10",
+                "market_cap": "132766738270.000000",
+                "tev": "147455738270.000000",
+                "shares_outstanding": 313099562,
+            },
+            {
+                "date": "2024-04-11",
+                "market_cap": "132416066761.000000",
+                "tev": "147105066761.000000",
+                "shares_outstanding": 313099562,
+            },
+        ],
+    }
+
+    def test_get_capitalization_from_identifiers(self, requests_mock: Mocker, mock_client: Client):
+        """
+        GIVEN the GetCapitalizationFromIdentifiers tool
+        WHEN we request the SPGI market cap
+        THEN we get back the SPGI market cap
+        """
+        requests_mock.get(
+            url=f"https://kfinance.kensho.com/api/v1/market_cap/{SPGI_COMPANY_ID}/none/none",
+            json=self.market_caps_resp,
+        )
+
+        expected_response = {
+            "SPGI": [
+                {"date": "2024-04-10", "market_cap": {"unit": "USD", "value": "132766738270.00"}},
+                {"date": "2024-04-11", "market_cap": {"unit": "USD", "value": "132416066761.00"}},
+            ]
+        }
+
+        tool = GetCapitalizationFromIdentifiers(kfinance_client=mock_client)
+        args = GetCapitalizationFromIdentifiersArgs(
+            identifiers=["SPGI"], capitalization=Capitalization.market_cap
+        )
+        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 GetCapitalizationFromIdentifiers tool
+        WHEN we request most recent market caps for multiple companies
+        THEN we only get back the most recent market cap for each company
+        """
+        expected_response = {
+            "C_1": [
+                {"date": "2024-04-10", "market_cap": {"unit": "USD", "value": "132766738270.00"}}
+            ],
+            "C_2": [
+                {"date": "2024-04-10", "market_cap": {"unit": "USD", "value": "132766738270.00"}}
+            ],
+        }
+
+        company_ids = [1, 2]
+        for company_id in company_ids:
+            requests_mock.get(
+                url=f"https://kfinance.kensho.com/api/v1/market_cap/{company_id}/none/none",
+                json=self.market_caps_resp,
+            )
+        tool = GetCapitalizationFromIdentifiers(kfinance_client=mock_client)
+        args = GetCapitalizationFromIdentifiersArgs(
+            identifiers=[f"{COMPANY_ID_PREFIX}{company_id}" for company_id in company_ids],
+            capitalization=Capitalization.market_cap,
+        )
+
+        response = tool.run(args.model_dump(mode="json"))
+        assert response == expected_response

kfinance/domains/companies/company_identifiers.py

@@ -0,0 +1,175 @@
+from dataclasses import dataclass
+from typing import Hashable, Protocol
+
+from kfinance.client.batch_request_handling import Task, process_tasks_in_thread_pool_executor
+from kfinance.client.fetch import KFinanceApiClient
+from kfinance.domains.companies.company_models import COMPANY_ID_PREFIX, IdentificationTriple
+
+
+class CompanyIdentifier(Protocol, Hashable):
+    """A CompanyIdentifier is an identifier that can be resolved to a company, security, or trading item id.
+
+    The two current identifiers are:
+    - Ticker/CUSIP/ISIN (resolve through ID triple)
+    - company id (from tools like business relationships)
+
+    These identifiers both have ways of fetching company, security, and trading
+    item ids but the paths differ, so this protocol defines the functional requirements
+    and leaves the implementation to sub classes.
+    """
+
+    api_client: KFinanceApiClient
+
+    def fetch_company_id(self) -> int:
+        """Return the company_id associated with the CompanyIdentifier."""
+
+    def fetch_security_id(self) -> int:
+        """Return the security_id associated with the CompanyIdentifier."""
+
+    def fetch_trading_item_id(self) -> int:
+        """Return the trading_item_id associated with the CompanyIdentifier."""
+
+
+@dataclass
+class Identifier(CompanyIdentifier):
+    """An identifier (ticker, CUSIP, ISIN), which can be resolved to company, security, or trading item id.
+
+    The resolution happens by fetching the id triple.
+    """
+
+    identifier: str
+    api_client: KFinanceApiClient
+    _id_triple: IdentificationTriple | None = None
+
+    def __str__(self) -> str:
+        return self.identifier
+
+    def __hash__(self) -> int:
+        return hash(self.identifier)
+
+    @property
+    def id_triple(self) -> IdentificationTriple:
+        """Return the id triple of the company."""
+        if self._id_triple is None:
+            id_triple_resp = self.api_client.fetch_id_triple(identifier=self.identifier)
+            self._id_triple = IdentificationTriple(
+                trading_item_id=id_triple_resp["trading_item_id"],
+                security_id=id_triple_resp["security_id"],
+                company_id=id_triple_resp["company_id"],
+            )
+        return self._id_triple
+
+    def fetch_company_id(self) -> int:
+        """Return the company_id associated with the Identifier."""
+        return self.id_triple.company_id
+
+    def fetch_security_id(self) -> int:
+        """Return the security_id associated with the Identifier."""
+        return self.id_triple.security_id
+
+    def fetch_trading_item_id(self) -> int:
+        """Return the trading_item_id associated with the Identifier."""
+        return self.id_triple.trading_item_id
+
+
+@dataclass
+class CompanyId(CompanyIdentifier):
+    """A company id, which can be resolved to security and trading item id.
+
+    The resolution happens by fetching the primary security and trading item id
+    associated with the company id.
+    """
+
+    company_id: int
+    api_client: KFinanceApiClient
+    _security_id: int | None = None
+    _trading_item_id: int | None = None
+
+    def __str__(self) -> str:
+        return f"{COMPANY_ID_PREFIX}{self.company_id}"
+
+    def __hash__(self) -> int:
+        return hash(self.company_id)
+
+    def fetch_company_id(self) -> int:
+        """Return the company_id."""
+        return self.company_id
+
+    def fetch_security_id(self) -> int:
+        """Return the security_id associated with the CompanyId."""
+        if self._security_id is None:
+            security_resp = self.api_client.fetch_primary_security(company_id=self.company_id)
+            self._security_id = security_resp["primary_security"]
+        return self._security_id
+
+    def fetch_trading_item_id(self) -> int:
+        """Return the trading_item_id associated with the CompanyId."""
+        if self._trading_item_id is None:
+            trading_item_resp = self.api_client.fetch_primary_trading_item(
+                security_id=self.fetch_security_id()
+            )
+            self._trading_item_id = trading_item_resp["primary_trading_item"]
+        return self._trading_item_id
+
+
+def parse_identifiers(
+    identifiers: list[str], api_client: KFinanceApiClient
+) -> list[CompanyIdentifier]:
+    """Return a list of CompanyIdentifier based on a list of string identifiers."""
+
+    parsed_identifiers: list[CompanyIdentifier] = []
+    for identifier in identifiers:
+        if identifier.startswith(COMPANY_ID_PREFIX):
+            parsed_identifiers.append(
+                CompanyId(
+                    company_id=int(identifier[len(COMPANY_ID_PREFIX) :]), api_client=api_client
+                )
+            )
+        else:
+            parsed_identifiers.append(Identifier(identifier=identifier, api_client=api_client))
+
+    return parsed_identifiers
+
+
+def fetch_company_ids_from_identifiers(
+    identifiers: list[CompanyIdentifier], api_client: KFinanceApiClient
+) -> dict[CompanyIdentifier, int]:
+    """Resolve a list of CompanyIdentifier to the corresponding company_ids."""
+
+    tasks = [
+        Task(
+            func=identifier.fetch_company_id,
+            result_key=identifier,
+        )
+        for identifier in identifiers
+    ]
+    return process_tasks_in_thread_pool_executor(api_client=api_client, tasks=tasks)
+
+
+def fetch_security_ids_from_identifiers(
+    identifiers: list[CompanyIdentifier], api_client: KFinanceApiClient
+) -> dict[CompanyIdentifier, int]:
+    """Resolve a list of CompanyIdentifier to the corresponding security_ids."""
+
+    tasks = [
+        Task(
+            func=identifier.fetch_security_id,
+            result_key=identifier,
+        )
+        for identifier in identifiers
+    ]
+    return process_tasks_in_thread_pool_executor(api_client=api_client, tasks=tasks)
+
+
+def fetch_trading_item_ids_from_identifiers(
+    identifiers: list[CompanyIdentifier], api_client: KFinanceApiClient
+) -> dict[CompanyIdentifier, int]:
+    """Resolve a list of CompanyIdentifier to the corresponding trading_item_ids."""
+    tasks = [
+        Task(
+            func=identifier.fetch_trading_item_id,
+            result_key=identifier,
+        )
+        for identifier in identifiers
+    ]
+    return process_tasks_in_thread_pool_executor(api_client=api_client, tasks=tasks)

kfinance/domains/companies/company_models.py

@@ -0,0 +1,27 @@
+from typing import NamedTuple
+
+from pydantic import BaseModel, field_serializer
+
+
+COMPANY_ID_PREFIX = "C_"
+
+
+class CompanyIdAndName(BaseModel):
+    """A company_id and name"""
+
+    company_id: int
+    company_name: str
+
+    @field_serializer("company_id")
+    def serialize_with_prefix(self, company_id: int) -> str:
+        """Serialize the company_id with a prefix ("C_<company_id>").
+
+        Including the prefix allows us to distinguish tickers and company_ids.
+        """
+        return f"{COMPANY_ID_PREFIX}{company_id}"
+
+
+class IdentificationTriple(NamedTuple):
+    trading_item_id: int
+    security_id: int
+    company_id: int

kfinance/domains/companies/company_tools.py

@@ -0,0 +1,66 @@
+from textwrap import dedent
+from typing import Type
+
+from pydantic import BaseModel
+
+from kfinance.client.batch_request_handling import Task, process_tasks_in_thread_pool_executor
+from kfinance.client.permission_models import Permission
+from kfinance.domains.companies.company_identifiers import (
+    fetch_company_ids_from_identifiers,
+    parse_identifiers,
+)
+from kfinance.integrations.tool_calling.tool_calling_models import (
+    KfinanceTool,
+    ToolArgsWithIdentifiers,
+)
+
+
+class GetInfoFromIdentifiers(KfinanceTool):
+    name: str = "get_info_from_identifiers"
+    description: str = dedent("""
+        Get the information associated with a list of identifiers. Info includes company name, status, type, simple industry, number of employees (if available), founding date, webpage, HQ address, HQ city, HQ zip code, HQ state, HQ country, and HQ country iso code.
+
+        - When possible, pass multiple identifiers in a single call rather than making multiple calls.
+    """).strip()
+    args_schema: Type[BaseModel] = ToolArgsWithIdentifiers
+    accepted_permissions: set[Permission] | None = None
+
+    def _run(self, identifiers: list[str]) -> dict:
+        """Sample response:
+
+        {
+            "SPGI": {
+                "name": "S&P Global Inc.",
+                "status": "Operating",
+                "type": "Public Company",
+                "simple_industry": "Capital Markets",
+                "number_of_employees": "42350.0000",
+                "founding_date": "1860-01-01",
+                "webpage": "www.spglobal.com",
+                "address": "55 Water Street",
+                "city": "New York",
+                "zip_code": "10041-0001",
+                "state": "New York",
+                "country": "United States",
+                "iso_country": "USA"
+            }
+        }
+        """
+        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_info,
+                kwargs=dict(company_id=company_id),
+                result_key=identifier,
+            )
+            for identifier, company_id in identifiers_to_company_ids.items()
+        ]
+
+        info_responses = process_tasks_in_thread_pool_executor(api_client=api_client, tasks=tasks)
+
+        return {str(identifier): result for identifier, result in info_responses.items()}

kfinance/domains/companies/tests/test_company_tools.py

@@ -0,0 +1,26 @@
+from requests_mock import Mocker
+
+from kfinance.client.kfinance import Client
+from kfinance.conftest import SPGI_COMPANY_ID
+from kfinance.domains.companies.company_tools import GetInfoFromIdentifiers
+from kfinance.integrations.tool_calling.tool_calling_models import ToolArgsWithIdentifiers
+
+
+class TestGetInfoFromIdentifiers:
+    def test_get_info_from_identifiers(self, mock_client: Client, requests_mock: Mocker):
+        """
+        GIVEN the GetInfoFromIdentifiers tool
+        WHEN request info for SPGI
+        THEN we get back info for SPGI
+        """
+
+        info_resp = {"name": "S&P Global Inc.", "status": "Operating"}
+        expected_response = {"SPGI": info_resp}
+        requests_mock.get(
+            url=f"https://kfinance.kensho.com/api/v1/info/{SPGI_COMPANY_ID}",
+            json=info_resp,
+        )
+
+        tool = GetInfoFromIdentifiers(kfinance_client=mock_client)
+        resp = tool.run(ToolArgsWithIdentifiers(identifiers=["SPGI"]).model_dump(mode="json"))
+        assert resp == expected_response

kfinance/{models → domains/competitors}/competitor_models.py

@@ -1,5 +1,8 @@
+from pydantic import BaseModel
 from strenum import StrEnum
 
+from kfinance.domains.companies.company_models import CompanyIdAndName
+
 
 class CompetitorSource(StrEnum):
     """The source type of the competitor information: 'filing' (from SEC filings), 'key_dev' (from key developments), 'contact' (from contact relationships), 'third_party' (from third-party sources), 'self_identified' (self-identified), 'named_by_competitor' (from competitor's perspective)."""
@@ -11,3 +14,7 @@ class CompetitorSource(StrEnum):
     third_party = "third_party"
     self_identified = "self_identified"
     named_by_competitor = "named_by_competitor"
+
+
+class CompetitorResponse(BaseModel):
+    competitors: list[CompanyIdAndName]

kfinance/domains/competitors/competitor_tools.py

@@ -0,0 +1,62 @@
+from kfinance.client.batch_request_handling import Task, process_tasks_in_thread_pool_executor
+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.competitors.competitor_models import CompetitorResponse, CompetitorSource
+from kfinance.integrations.tool_calling.tool_calling_models import (
+    KfinanceTool,
+    ToolArgsWithIdentifiers,
+)
+
+
+class GetCompetitorsFromIdentifiersArgs(ToolArgsWithIdentifiers):
+    # no description because the description for enum fields comes from the enum docstring.
+    competitor_source: CompetitorSource
+
+
+class GetCompetitorsFromIdentifiers(KfinanceTool):
+    name: str = "get_competitors_from_identifiers"
+    description: str = "Retrieves a list of company_id and company_name that are competitors for a list of companies, optionally filtered by the source of the competitor information."
+    args_schema = GetCompetitorsFromIdentifiersArgs
+    accepted_permissions: set[Permission] | None = {Permission.CompetitorsPermission}
+
+    def _run(
+        self,
+        identifiers: list[str],
+        competitor_source: CompetitorSource,
+    ) -> dict:
+        """Sample response:
+
+        {
+            SPGI: {
+                {'company_id': "C_35352", 'company_name': 'The Descartes Systems Group Inc.'},
+                {'company_id': "C_4003514", 'company_name': 'London Stock Exchange Group plc'}
+            }
+        }
+        """
+
+        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_competitors,
+                kwargs=dict(company_id=company_id, competitor_source=competitor_source),
+                result_key=identifier,
+            )
+            for identifier, company_id in identifiers_to_company_ids.items()
+        ]
+
+        competitor_responses: dict[Identifier, CompetitorResponse] = (
+            process_tasks_in_thread_pool_executor(api_client=api_client, tasks=tasks)
+        )
+        return {
+            str(identifier): competitors.model_dump(mode="json")
+            for identifier, competitors in competitor_responses.items()
+        }

kfinance/domains/competitors/tests/test_competitor_tools.py

@@ -0,0 +1,45 @@
+from requests_mock import Mocker
+
+from kfinance.client.kfinance import Client
+from kfinance.conftest import SPGI_COMPANY_ID
+from kfinance.domains.competitors.competitor_models import CompetitorSource
+from kfinance.domains.competitors.competitor_tools import (
+    GetCompetitorsFromIdentifiers,
+    GetCompetitorsFromIdentifiersArgs,
+)
+
+
+class TestGetCompetitorsFromIdentifiers:
+    def test_get_competitors_from_identifiers(self, mock_client: Client, requests_mock: Mocker):
+        """
+        GIVEN the GetCompetitorsFromIdentifiers tool
+        WHEN we request the SPGI competitors that are named by competitors
+        THEN we get back the SPGI competitors that are named by competitors
+        """
+        competitors_response = {
+            "competitors": [
+                {"company_id": 35352, "company_name": "The Descartes Systems Group Inc."},
+                {"company_id": 4003514, "company_name": "London Stock Exchange Group plc"},
+            ]
+        }
+        expected_response = {
+            "SPGI": {
+                "competitors": [
+                    {"company_id": "C_35352", "company_name": "The Descartes Systems Group Inc."},
+                    {"company_id": "C_4003514", "company_name": "London Stock Exchange Group plc"},
+                ]
+            }
+        }
+
+        requests_mock.get(
+            url=f"https://kfinance.kensho.com/api/v1/competitors/{SPGI_COMPANY_ID}/named_by_competitor",
+            # truncated from the original API response
+            json=competitors_response,
+        )
+
+        tool = GetCompetitorsFromIdentifiers(kfinance_client=mock_client)
+        args = GetCompetitorsFromIdentifiersArgs(
+            identifiers=["SPGI"], competitor_source=CompetitorSource.named_by_competitor
+        )
+        response = tool.run(args.model_dump(mode="json"))
+        assert response == expected_response