ask-candid / ask_candid /tools /elastic /index_data_tool.py
brainsqueeze's picture
Upload 36 files
a0e37e2 verified
from typing import Type, Optional
import logging
from pydantic import BaseModel, Field
from elasticsearch import Elasticsearch
from langchain.callbacks.manager import CallbackManagerForToolRun
from langchain.tools.base import BaseTool
from ask_candid.base.config.connections import SEMANTIC_ELASTIC_QA
logging.basicConfig(level="INFO")
logger = logging.getLogger("elasticsearch_playground")
es = Elasticsearch(
cloud_id=SEMANTIC_ELASTIC_QA.cloud_id,
api_key=SEMANTIC_ELASTIC_QA.api_key,
verify_certs=True,
request_timeout=60 * 3
)
class IndexShowDataInput(BaseModel):
"""Input for the index show data tool."""
index_name: str = Field(
..., description="The name of the index for which the data is to be retrieved"
)
class IndexShowDataTool(BaseTool):
"""Tool for getting a list of entries from an ElasticSearch index, helpful to figure out what data is available."""
name: str = "elastic_index_show_data" # Added type annotation
description: str = (
"Input is an index name, output is a JSON based string with an extract of the data of the index"
)
args_schema: Optional[Type[BaseModel]] = (
IndexShowDataInput # This should be placed before methods
)
def _run(
self,
index_name: str,
run_manager: Optional[CallbackManagerForToolRun] = None,
) -> str:
"""Get all indices in the Elasticsearch server, usually separated by a line break."""
try:
# Ensure `es` is properly initialized before this method is called
res = es.search(
index=index_name,
from_=0,
size=20,
query={"match_all": {}},
)
return str(res["hits"])
except Exception as e:
print(e)
logger.exception("Could not fetch index data for %s", index_name)
return ""