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

kfinance/domains/capitalizations/capitalization_tools.py

@@ -18,14 +18,17 @@ 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
+        description="The start date for historical capitalization retrieval. Use null for latest values. For annual data, use January 1st of the year.",
+        default=None,
     )
     end_date: date | None = Field(
-        description="The end date for historical capitalization retrieval", default=None
+        description="The end date for historical capitalization retrieval. Use null for latest values. For annual data, use December 31st of the year.",
+        default=None,
     )
 
 
 class GetCapitalizationFromIdentifiersResp(ToolRespWithErrors):
+    capitalization: Capitalization
     results: dict[str, Capitalizations]
 
 
@@ -35,11 +38,20 @@ class GetCapitalizationFromIdentifiers(KfinanceTool):
         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.
+        - When requesting the most recent values, leave start_date and end_date null.
+        - For annual data (e.g., "market cap in 2020", "FY2021 values"), use the full year range: start_date as January 1st and end_date as December 31st.
+        - For "latest" or "current" values, always leave dates null to get the most recent data point.
+        - Only specify date ranges when the user explicitly requests historical data over a specific period.
+
+        Examples:
+        Query: "What are the market caps of Visa and Mastercard?"
+        Function: get_capitalization_from_identifiers(capitalization="market_cap", identifiers=["Visa", "Mastercard"], start_date=null, end_date=null)
+
+        Query: "What was MDT's market cap in 2020?"
+        Function: get_capitalization_from_identifiers(capitalization="market_cap", identifiers=["MDT"], start_date="2020-01-01", end_date="2020-12-31")
 
-        Example:
-        Query: "What are the market caps of AAPL and WMT?"
-        Function: get_capitalization_from_identifiers(capitalization=Capitalization.market_cap, identifiers=["AAPL", "WMT"])
+        Query: "Market cap trends for MSFT from Q1 2020 to Q3 2021"
+        Function: get_capitalization_from_identifiers(capitalization="market_cap", identifiers=["MSFT"], start_date="2020-01-01", end_date="2021-09-30")
     """).strip()
     args_schema: Type[BaseModel] = GetCapitalizationFromIdentifiersArgs
     accepted_permissions: set[Permission] | None = {Permission.PricingPermission}
@@ -50,13 +62,13 @@ class GetCapitalizationFromIdentifiers(KfinanceTool):
         capitalization: Capitalization,
         start_date: str | None = None,
         end_date: str | None = None,
-    ) -> dict:
+    ) -> GetCapitalizationFromIdentifiersResp:
         """Sample response:
 
         {
+            'capitalization': 'market_cap'
             'results': {
                 'SPGI': {
-                    'capitalizations': [
                         {'date': '2024-04-10', 'market_cap': {'value': '132766738270.00', 'unit': 'USD'}},
                         {'date': '2024-04-11', 'market_cap': {'value': '132416066761.00', 'unit': 'USD'}}
                     ]
@@ -83,7 +95,11 @@ class GetCapitalizationFromIdentifiers(KfinanceTool):
             api_client=api_client, tasks=tasks
         )
 
-        for capitalization_response in capitalization_responses.values():
+        for identifier, capitalization_response in capitalization_responses.items():
+            # If we get an empty response for a company, assign an empty object
+            if not capitalization_response:
+                capitalization_responses[identifier] = Capitalizations(capitalizations=list())
+                capitalization_response = capitalization_responses[identifier]
             # If we return results for more than one company and the start and end dates are unset,
             # truncate data to only return the most recent datapoint.
             if len(capitalization_responses) > 1 and start_date is None and end_date is None:
@@ -92,15 +108,17 @@ class GetCapitalizationFromIdentifiers(KfinanceTool):
                 ]
             # Set capitalizations that were not requested to None.
             # That way, they can be skipped for serialization via `exclude_none=True`
-            for daily_capitalization in capitalization_response.capitalizations:
-                if capitalization is not Capitalization.market_cap:
-                    daily_capitalization.market_cap = None
-                if capitalization is not Capitalization.tev:
-                    daily_capitalization.tev = None
-                if capitalization is not Capitalization.shares_outstanding:
-                    daily_capitalization.shares_outstanding = None
-
-        resp_model = GetCapitalizationFromIdentifiersResp(
-            results=capitalization_responses, errors=list(id_triple_resp.errors.values())
+            if capitalization_response.capitalizations:
+                for daily_capitalization in capitalization_response.capitalizations:
+                    if capitalization is not Capitalization.market_cap:
+                        daily_capitalization.market_cap = None
+                    if capitalization is not Capitalization.tev:
+                        daily_capitalization.tev = None
+                    if capitalization is not Capitalization.shares_outstanding:
+                        daily_capitalization.shares_outstanding = None
+
+        return GetCapitalizationFromIdentifiersResp(
+            capitalization=capitalization,
+            results=capitalization_responses,
+            errors=list(id_triple_resp.errors.values()),
         )
-        return resp_model.model_dump(mode="json", exclude_none=True)

kfinance/domains/capitalizations/tests/test_capitalization_tools.py

@@ -1,11 +1,20 @@
+from datetime import date
+from decimal import Decimal
+
 from requests_mock import Mocker
 
 from kfinance.client.kfinance import Client
+from kfinance.client.models.decimal_with_unit import Money
 from kfinance.conftest import SPGI_COMPANY_ID
-from kfinance.domains.capitalizations.capitalization_models import Capitalization
+from kfinance.domains.capitalizations.capitalization_models import (
+    Capitalization,
+    Capitalizations,
+    DailyCapitalization,
+)
 from kfinance.domains.capitalizations.capitalization_tools import (
     GetCapitalizationFromIdentifiers,
     GetCapitalizationFromIdentifiersArgs,
+    GetCapitalizationFromIdentifiersResp,
 )
 from kfinance.domains.companies.company_models import COMPANY_ID_PREFIX
 
@@ -40,25 +49,30 @@ class TestGetCapitalizationFromCompanyIds:
             json=self.market_caps_resp,
         )
 
-        expected_response = {
-            "results": {
-                "SPGI": {
-                    "capitalizations": [
-                        {
-                            "date": "2024-04-10",
-                            "market_cap": {"value": "132766738270.00", "unit": "USD"},
-                        },
-                        {
-                            "date": "2024-04-11",
-                            "market_cap": {"value": "132416066761.00", "unit": "USD"},
-                        },
+        expected_response = GetCapitalizationFromIdentifiersResp(
+            capitalization=Capitalization.market_cap,
+            results={
+                "SPGI": Capitalizations(
+                    market_caps=[
+                        DailyCapitalization(
+                            date=date(2024, 4, 10),
+                            market_cap=Money(value=Decimal(132766738270), unit="USD"),
+                            tev=None,
+                            shares_outstanding=None,
+                        ),
+                        DailyCapitalization(
+                            date=date(2024, 4, 11),
+                            market_cap=Money(value=Decimal(132416066761), unit="USD"),
+                            tev=None,
+                            shares_outstanding=None,
+                        ),
                     ]
-                }
+                ),
             },
-            "errors": [
+            errors=[
                 "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
             ],
-        }
+        )
 
         tool = GetCapitalizationFromIdentifiers(kfinance_client=mock_client)
         args = GetCapitalizationFromIdentifiersArgs(
@@ -67,32 +81,48 @@ class TestGetCapitalizationFromCompanyIds:
         response = tool.run(args.model_dump(mode="json"))
         assert response == expected_response
 
+    def test_get_capitalization_from_identifiers_property_400(
+        self, requests_mock: Mocker, mock_client: Client
+    ) -> None:
+        requests_mock.get(
+            url=f"https://kfinance.kensho.com/api/v1/market_cap/1/none/none",
+            status_code=400,
+        )
+
+        expected_response = GetCapitalizationFromIdentifiersResp(
+            capitalization=Capitalization.market_cap,
+            results={"C_1": Capitalizations(market_caps=list())},
+        )
+
+        tool = GetCapitalizationFromIdentifiers(kfinance_client=mock_client)
+        args = GetCapitalizationFromIdentifiersArgs(
+            identifiers=["C_1"], 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 = {
-            "results": {
-                "C_1": {
-                    "capitalizations": [
-                        {
-                            "date": "2024-04-11",
-                            "market_cap": {"unit": "USD", "value": "132416066761.00"},
-                        }
-                    ]
-                },
-                "C_2": {
-                    "capitalizations": [
-                        {
-                            "date": "2024-04-11",
-                            "market_cap": {"unit": "USD", "value": "132416066761.00"},
-                        }
-                    ]
-                },
-            }
-        }
+
+        capitalization = Capitalizations(
+            market_caps=[
+                DailyCapitalization(
+                    date=date(2024, 4, 11),
+                    market_cap=Money(value=Decimal(132416066761), unit="USD"),
+                    tev=None,
+                    shares_outstanding=None,
+                )
+            ]
+        )
+
+        expected_response = GetCapitalizationFromIdentifiersResp(
+            capitalization=Capitalization.market_cap,
+            results={"C_1": capitalization, "C_2": capitalization},
+        )
 
         company_ids = [1, 2]
         for company_id in company_ids:

kfinance/domains/companies/company_models.py

@@ -84,14 +84,20 @@ class UnifiedIdTripleResponse(BaseModel):
 
 
         """
-        output: dict[str, dict] = dict(identifiers_to_id_triples=dict(), errors=dict())
+        # Separate successful and failed resolutions for kfinance api responses
         if isinstance(data, dict) and "data" in data:
+            output: dict[str, dict] = dict(identifiers_to_id_triples=dict(), errors=dict())
+
             for key, val in data["data"].items():
                 if "error" in val:
                     output["errors"][key] = val["error"]
                 else:
                     output["identifiers_to_id_triples"][key] = val
-        return output
+            return output
+        # In all other cases (e.g. UnifiedIdTripleResponse directly initialized),
+        # just return the data.
+        else:
+            return data
 
     def filter_out_companies_without_security_ids(self) -> None:
         """Filter out companies that don't have a security_id and add an error for them."""
@@ -119,6 +125,20 @@ class UnifiedIdTripleResponse(BaseModel):
             )
             self.identifiers_to_id_triples.pop(identifier)
 
+    def get_identifier_from_company_id(self, company_id: int) -> str:
+        """Return the (originally passed) identifier from a company id."""
+        if not hasattr(self, "_company_id_to_identifier"):
+            self._company_id_to_identifier = {
+                id_triple.company_id: identifier
+                for identifier, id_triple in self.identifiers_to_id_triples.items()
+            }
+        return self._company_id_to_identifier[company_id]
+
+    @property
+    def company_ids(self) -> list[int]:
+        """Returns a list of all company ids in the response."""
+        return [id_triple.company_id for id_triple in self.identifiers_to_id_triples.values()]
+
 
 class CompanyDescriptions(BaseModel):
     """A company summary and description"""

kfinance/domains/companies/company_tools.py

@@ -5,7 +5,11 @@ 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_models import CompanyDescriptions, CompanyOtherNames
+from kfinance.domains.companies.company_models import (
+    COMPANY_ID_PREFIX,
+    CompanyDescriptions,
+    CompanyOtherNames,
+)
 from kfinance.integrations.tool_calling.tool_calling_models import (
     KfinanceTool,
     ToolArgsWithIdentifiers,
@@ -20,14 +24,21 @@ class GetInfoFromIdentifiersResp(ToolRespWithErrors):
 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.
+        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, HQ country iso code, and CIQ company_id.
 
         - When possible, pass multiple identifiers in a single call rather than making multiple calls.
+
+        Examples:
+        Query: "What's the company information for Northrop Grumman and Lockheed Martin?"
+        Function: get_info_from_identifiers(identifiers=["Northrop Grumman", "Lockheed Martin"])
+
+        Query: "Get company info for UBER and LYFT"
+        Function: get_info_from_identifiers(identifiers=["UBER", "LYFT"])
     """).strip()
     args_schema: Type[BaseModel] = ToolArgsWithIdentifiers
     accepted_permissions: set[Permission] | None = None
 
-    def _run(self, identifiers: list[str]) -> dict:
+    def _run(self, identifiers: list[str]) -> GetInfoFromIdentifiersResp:
         """Sample response:
 
         {   "results": {
@@ -44,7 +55,8 @@ class GetInfoFromIdentifiers(KfinanceTool):
                     "zip_code": "10041-0001",
                     "state": "New York",
                     "country": "United States",
-                    "iso_country": "USA"
+                    "iso_country": "USA",
+                    "company_id": "C_21719"
                 }
             },
             "errors": [['No identification triple found for the provided identifier: NON-EXISTENT of type: ticker']
@@ -65,10 +77,13 @@ class GetInfoFromIdentifiers(KfinanceTool):
         info_responses: dict[str, dict] = process_tasks_in_thread_pool_executor(
             api_client=api_client, tasks=tasks
         )
-        resp_model = GetInfoFromIdentifiersResp(
+
+        for identifier, id_triple in id_triple_resp.identifiers_to_id_triples.items():
+            info_responses[identifier]["company_id"] = f"{COMPANY_ID_PREFIX}{id_triple.company_id}"
+
+        return GetInfoFromIdentifiersResp(
             results=info_responses, errors=list(id_triple_resp.errors.values())
         )
-        return resp_model.model_dump(mode="json")
 
 
 class GetCompanyOtherNamesFromIdentifiersResp(ToolRespWithErrors):
@@ -81,6 +96,13 @@ class GetCompanyOtherNamesFromIdentifiers(KfinanceTool):
         Given a list of identifiers, fetch the alternate, historical, and native names associated with each identifier. Alternate names are additional names a company might go by (for example, Hewlett-Packard Company also goes by the name HP). Historical names are previous names for the company if it has changed over time. Native names are primary non-Latin character native names for global companies, including languages such as Arabic, Russian, Greek, Japanese, etc. This also includes limited history on native name changes.
 
         - When possible, pass multiple identifiers in a single call rather than making multiple calls.
+
+        Examples:
+        Query: "What are the alternate names for Meta and Alphabet?"
+        Function: get_company_other_names_from_identifiers(identifiers=["Meta", "Alphabet"])
+
+        Query: "Get other names for NSRGY"
+        Function: get_company_other_names_from_identifiers(identifiers=["NSRGY"])
     """).strip()
     args_schema: Type[BaseModel] = ToolArgsWithIdentifiers
     accepted_permissions: set[Permission] | None = {Permission.CompanyIntelligencePermission}
@@ -88,7 +110,7 @@ class GetCompanyOtherNamesFromIdentifiers(KfinanceTool):
     def _run(
         self,
         identifiers: list[str],
-    ) -> dict:
+    ) -> GetCompanyOtherNamesFromIdentifiersResp:
         api_client = self.kfinance_client.kfinance_api_client
         id_triple_resp = api_client.unified_fetch_id_triples(identifiers=identifiers)
         tasks = [
@@ -102,10 +124,9 @@ class GetCompanyOtherNamesFromIdentifiers(KfinanceTool):
         info_responses: dict[str, CompanyOtherNames] = process_tasks_in_thread_pool_executor(
             api_client=api_client, tasks=tasks
         )
-        resp_model = GetCompanyOtherNamesFromIdentifiersResp(
+        return GetCompanyOtherNamesFromIdentifiersResp(
             results=info_responses, errors=list(id_triple_resp.errors.values())
         )
-        return resp_model.model_dump(mode="json")
 
 
 class GetCompanySummaryFromIdentifiersResp(ToolRespWithErrors):
@@ -115,9 +136,16 @@ class GetCompanySummaryFromIdentifiersResp(ToolRespWithErrors):
 class GetCompanySummaryFromIdentifiers(KfinanceTool):
     name: str = "get_company_summary_from_identifiers"
     description: str = dedent("""
-        Get one paragraph summary/short descriptions of companies, including information about the company's primary business, products and services offered and their applications, business segment details, client/customer groups served, geographic markets served, distribution channels, strategic alliances/partnerships, founded/incorporated year, latest former name, and headquarters and additional offices."
+        Get one paragraph summary/short descriptions of companies, including information about the company's primary business, products and services offered and their applications, business segment details, client/customer groups served, geographic markets served, distribution channels, strategic alliances/partnerships, founded/incorporated year, latest former name, and headquarters and additional offices.
 
         - When possible, pass multiple identifiers in a single call rather than making multiple calls.
+
+        Examples:
+        Query: "Give me summaries of Tesla and General Motors"
+        Function: get_company_summary_from_identifiers(identifiers=["Tesla", "General Motors"])
+
+        Query: "What are the summaries for F and STLA?"
+        Function: get_company_summary_from_identifiers(identifiers=["F", "STLA"])
     """).strip()
     args_schema: Type[BaseModel] = ToolArgsWithIdentifiers
     accepted_permissions: set[Permission] | None = {Permission.CompanyIntelligencePermission}
@@ -125,7 +153,7 @@ class GetCompanySummaryFromIdentifiers(KfinanceTool):
     def _run(
         self,
         identifiers: list[str],
-    ) -> dict:
+    ) -> GetCompanySummaryFromIdentifiersResp:
         api_client = self.kfinance_client.kfinance_api_client
         id_triple_resp = api_client.unified_fetch_id_triples(identifiers=identifiers)
 
@@ -147,10 +175,9 @@ class GetCompanySummaryFromIdentifiers(KfinanceTool):
             for identifier, descriptions in company_description_responses.items()
         }
 
-        resp_model = GetCompanySummaryFromIdentifiersResp(
+        return GetCompanySummaryFromIdentifiersResp(
             results=summary_results, errors=list(id_triple_resp.errors.values())
         )
-        return resp_model.model_dump(mode="json")
 
 
 class GetCompanyDescriptionFromIdentifiersResp(ToolRespWithErrors):
@@ -163,6 +190,13 @@ class GetCompanyDescriptionFromIdentifiers(KfinanceTool):
         Get detailed descriptions of companies, broken down into sections, which may include information about the company's Primary business, Segments (including Products and Services for each), Competition, Significant events, and History. Within the text, four spaces represent a new paragraph. Note that the description is divided into sections with headers, where each section has a new paragraph (four spaces) before and after the section header.
 
         - When possible, pass multiple identifiers in a single call rather than making multiple calls.
+
+        Examples:
+        Query: "Get detailed descriptions for Netflix and Disney"
+        Function: get_company_description_from_identifiers(identifiers=["Netflix", "Disney"])
+
+        Query: "What are the detailed company descriptions for KO and PEP?"
+        Function: get_company_description_from_identifiers(identifiers=["KO", "PEP"])
     """).strip()
     args_schema: Type[BaseModel] = ToolArgsWithIdentifiers
     accepted_permissions: set[Permission] | None = {Permission.CompanyIntelligencePermission}
@@ -170,7 +204,7 @@ class GetCompanyDescriptionFromIdentifiers(KfinanceTool):
     def _run(
         self,
         identifiers: list[str],
-    ) -> dict:
+    ) -> GetCompanyDescriptionFromIdentifiersResp:
         api_client = self.kfinance_client.kfinance_api_client
         id_triple_resp = api_client.unified_fetch_id_triples(identifiers=identifiers)
 
@@ -192,7 +226,6 @@ class GetCompanyDescriptionFromIdentifiers(KfinanceTool):
             for identifier, descriptions in company_description_responses.items()
         }
 
-        resp_model = GetCompanyDescriptionFromIdentifiersResp(
+        return GetCompanyDescriptionFromIdentifiersResp(
             results=description_results, errors=list(id_triple_resp.errors.values())
         )
-        return resp_model.model_dump(mode="json")

kfinance/domains/companies/tests/test_company_tools.py

@@ -2,11 +2,16 @@ 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.companies.company_tools import (
     GetCompanyDescriptionFromIdentifiers,
+    GetCompanyDescriptionFromIdentifiersResp,
     GetCompanyOtherNamesFromIdentifiers,
+    GetCompanyOtherNamesFromIdentifiersResp,
     GetCompanySummaryFromIdentifiers,
+    GetCompanySummaryFromIdentifiersResp,
     GetInfoFromIdentifiers,
+    GetInfoFromIdentifiersResp,
 )
 from kfinance.integrations.tool_calling.tool_calling_models import ToolArgsWithIdentifiers
 
@@ -19,13 +24,20 @@ class TestGetInfoFromIdentifiers:
         THEN we get back info for SPGI and an error for the non-existent company
         """
 
-        info_resp = {"name": "S&P Global Inc.", "status": "Operating"}
-        expected_response = {
-            "results": {"SPGI": info_resp},
-            "errors": [
-                "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
-            ],
+        info_resp = {
+            "name": "S&P Global Inc.",
+            "status": "Operating",
+            "company_id": f"{COMPANY_ID_PREFIX}{SPGI_COMPANY_ID}",
         }
+        expected_response = GetInfoFromIdentifiersResp.model_validate(
+            {
+                "results": {"SPGI": info_resp},
+                "errors": [
+                    "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
+                ],
+            }
+        )
+        del info_resp["company_id"]
         requests_mock.get(
             url=f"https://kfinance.kensho.com/api/v1/info/{SPGI_COMPANY_ID}",
             json=info_resp,
@@ -61,7 +73,9 @@ class TestGetCompanyDescriptions:
         tool = GetCompanySummaryFromIdentifiers(kfinance_client=mock_client)
         args = ToolArgsWithIdentifiers(identifiers=["SPGI"])
         response = tool.run(args.model_dump(mode="json"))
-        expected_response = {"results": {"SPGI": self.summary}}
+        expected_response = GetCompanySummaryFromIdentifiersResp.model_validate(
+            {"results": {"SPGI": self.summary}}
+        )
         assert response == expected_response
 
     def test_get_company_description_from_identifier(
@@ -81,7 +95,9 @@ class TestGetCompanyDescriptions:
         tool = GetCompanyDescriptionFromIdentifiers(kfinance_client=mock_client)
         args = ToolArgsWithIdentifiers(identifiers=["SPGI"])
         response = tool.run(args.model_dump(mode="json"))
-        expected_response = {"results": {"SPGI": self.description}}
+        expected_response = GetCompanyDescriptionFromIdentifiersResp.model_validate(
+            {"results": {"SPGI": self.description}}
+        )
         assert response == expected_response
 
 
@@ -121,5 +137,7 @@ class TestGetCompanyOtherNames:
         tool = GetCompanyOtherNamesFromIdentifiers(kfinance_client=mock_client)
         args = ToolArgsWithIdentifiers(identifiers=["SPGI"])
         response = tool.run(args.model_dump(mode="json"))
-        expected_response = {"results": {"SPGI": self.company_other_names_info}}
+        expected_response = GetCompanyOtherNamesFromIdentifiersResp.model_validate(
+            {"results": {"SPGI": self.company_other_names_info}}
+        )
         assert response == expected_response

kfinance/domains/competitors/competitor_tools.py

@@ -1,3 +1,6 @@
+from textwrap import dedent
+from typing import Type
+
 from kfinance.client.batch_request_handling import Task, process_tasks_in_thread_pool_executor
 from kfinance.client.permission_models import Permission
 from kfinance.domains.competitors.competitor_models import CompetitorResponse, CompetitorSource
@@ -19,15 +22,27 @@ class GetCompetitorsFromIdentifiersResp(ToolRespWithErrors):
 
 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
+    description: str = dedent("""
+        Retrieves a list of company_id and company_name that are competitors for a list of companies, filtered by the source of the competitor information.
+
+        - When possible, pass multiple identifiers in a single call rather than making multiple calls.
+        - Available competitor sources: all, 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)
+
+        Examples:
+        Query: "Who are Microsoft's competitors from SEC filings?"
+        Function: get_competitors_from_identifiers(identifiers=["Microsoft"], competitor_source="filing")
+
+        Query: "Get all competitors of AAPL and GOOGL"
+        Function: get_competitors_from_identifiers(identifiers=["AAPL", "GOOGL"], competitor_source="all")
+    """).strip()
+    args_schema: Type[GetCompetitorsFromIdentifiersArgs] = GetCompetitorsFromIdentifiersArgs
     accepted_permissions: set[Permission] | None = {Permission.CompetitorsPermission}
 
     def _run(
         self,
         identifiers: list[str],
         competitor_source: CompetitorSource,
-    ) -> dict:
+    ) -> GetCompetitorsFromIdentifiersResp:
         """Sample response:
 
         {
@@ -56,7 +71,6 @@ class GetCompetitorsFromIdentifiers(KfinanceTool):
         competitor_responses: dict[str, CompetitorResponse] = process_tasks_in_thread_pool_executor(
             api_client=api_client, tasks=tasks
         )
-        resp_model = GetCompetitorsFromIdentifiersResp(
+        return GetCompetitorsFromIdentifiersResp(
             results=competitor_responses, errors=list(id_triple_resp.errors.values())
         )
-        return resp_model.model_dump(mode="json")

kfinance/domains/competitors/tests/test_competitor_tools.py

@@ -6,6 +6,7 @@ from kfinance.domains.competitors.competitor_models import CompetitorSource
 from kfinance.domains.competitors.competitor_tools import (
     GetCompetitorsFromIdentifiers,
     GetCompetitorsFromIdentifiersArgs,
+    GetCompetitorsFromIdentifiersResp,
 )
 
 
@@ -24,25 +25,27 @@ class TestGetCompetitorsFromIdentifiers:
                 {"company_id": 4003514, "company_name": "London Stock Exchange Group plc"},
             ]
         }
-        expected_response = {
-            "results": {
-                "SPGI": {
-                    "competitors": [
-                        {
-                            "company_id": "C_35352",
-                            "company_name": "The Descartes Systems Group Inc.",
-                        },
-                        {
-                            "company_id": "C_4003514",
-                            "company_name": "London Stock Exchange Group plc",
-                        },
-                    ]
-                }
-            },
-            "errors": [
-                "No identification triple found for the provided identifier: NON-EXISTENT of type: ticker"
-            ],
-        }
+        expected_response = GetCompetitorsFromIdentifiersResp.model_validate(
+            {
+                "results": {
+                    "SPGI": {
+                        "competitors": [
+                            {
+                                "company_id": 35352,
+                                "company_name": "The Descartes Systems Group Inc.",
+                            },
+                            {
+                                "company_id": 4003514,
+                                "company_name": "London Stock Exchange Group plc",
+                            },
+                        ]
+                    }
+                },
+                "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/competitors/{SPGI_COMPANY_ID}/named_by_competitor",