codekingpro/portable-devtools

1from __future__ import annotations2 3import logging4from typing import Optional, Type5 6from langchain_core.callbacks import CallbackManagerForToolRun7from pydantic import BaseModel, Field, HttpUrl8 9from langchain_community.tools.edenai.edenai_base_tool import EdenaiTool10 11logger = logging.getLogger(__name__)12 13 14class IDParsingInput(BaseModel):15    query: HttpUrl = Field(description="url of the document to parse")16 17 18class EdenAiParsingIDTool(EdenaiTool):19    """Tool that queries the Eden AI  Identity parsing API.20 21    for api reference check edenai documentation:22    https://docs.edenai.co/reference/ocr_identity_parser_create.23 24    To use, you should have25    the environment variable ``EDENAI_API_KEY`` set with your API token.26    You can find your token here: https://app.edenai.run/admin/account/settings27 28    """29 30    name: str = "edenai_identity_parsing"31 32    description: str = (33        "A wrapper around edenai Services Identity parsing. "34        "Useful for when you have to extract information from an ID Document "35        "Input should be the string url of the document to parse."36    )37    args_schema: Type[BaseModel] = IDParsingInput38 39    feature: str = "ocr"40    subfeature: str = "identity_parser"41 42    language: Optional[str] = None43    """44    language of the text passed to the model.45    """46 47    def _parse_response(self, response: list) -> str:48        formatted_list: list = []49 50        if len(response) == 1:51            self._parse_json_multilevel(52                response[0]["extracted_data"][0], formatted_list53            )54        else:55            for entry in response:56                if entry.get("provider") == "eden-ai":57                    self._parse_json_multilevel(58                        entry["extracted_data"][0], formatted_list59                    )60 61        return "\n".join(formatted_list)62 63    def _run(64        self,65        query: str,66        run_manager: Optional[CallbackManagerForToolRun] = None,67    ) -> str:68        """Use the tool."""69        query_params = {70            "file_url": query,71            "language": self.language,72            "attributes_as_list": False,73        }74 75        return self._call_eden_ai(query_params)76