Delete ask_candid/agents

ask_candid/agents/elastic.py

@@ -1,842 +0,0 @@
-from typing import TypedDict, List
-from functools import partial
-import json
-import ast
-from ask_candid.base.api_base import BaseAPI
-import os
-import pandas as pd
-from pydantic import BaseModel, Field
-
-from langchain_core.runnables import RunnableSequence
-from langchain_core.language_models.llms import LLM
-from langchain.agents.openai_functions_agent.base import create_openai_functions_agent
-from langchain.agents.agent import AgentExecutor
-from langchain.agents.agent_types import AgentType
-from langchain.prompts import ChatPromptTemplate, PromptTemplate, MessagesPlaceholder
-from langchain.output_parsers import PydanticOutputParser
-from langchain.schema import BaseMessage
-from langchain.agents import create_tool_calling_agent, AgentExecutor
-from langchain_core.tools import Tool
-
-from langgraph.graph import StateGraph, END
-
-from ask_candid.tools.elastic.index_data_tool import IndexShowDataTool
-from ask_candid.tools.elastic.index_details_tool import IndexDetailsTool
-from ask_candid.tools.elastic.index_search_tool import create_search_tool
-
-tools = [
-    IndexShowDataTool(),
-    IndexDetailsTool(),
-    create_search_tool(pcs_codes={}),
-]
-
-
-class AutocodingAPI(BaseAPI):
-    def __init__(self):
-        super().__init__(
-            url=os.getenv("AUTOCODING_API_URL"),
-            headers={
-                "x-api-key": os.getenv("AUTOCODING_API_KEY"),
-                "Content-Type": "application/json",
-            },
-        )
-
-    def __call__(self, text: str, taxonomy: str = "pcs-v3"):
-        params = {"text": text, "taxonomy": taxonomy}
-        return self.get(**params)
-
-
-def find_subject_levels(filtered_df, subject_level_i, target_value):
-    """
-    Filters the DataFrame from the last valid NaN in 'Subject Level i' and retrieves corresponding values for lower levels.
-
-    Parameters:
-        filtered_df (pd.DataFrame): The input DataFrame.
-        subject_level_i (int): The subject level to filter from (1 to 4).
-        target_value (str): The value to search for in 'Subject Level i'.
-
-    Returns:
-        dict: A dictionary containing values for 'Subject Level i' to 'Subject Level 1'.
-        pd.DataFrame: The filtered DataFrame from the determined start index to the target_value row.
-    """
-    if subject_level_i < 1 or subject_level_i > 4:
-        raise ValueError("subject_level_i should be between 1 and 4")
-
-    # Define the target column dynamically
-    target_column = f"Subject Level {subject_level_i}"
-
-    # Find indices where the target column has the target value
-    target_indices = filtered_df[
-        filtered_df[target_column].astype(str).str.strip() == target_value
-    ].index
-
-    if target_indices.empty:
-        return {}, pd.DataFrame()  # Return empty if target_value is not found
-
-    # Get the first occurrence of the target value
-    first_target_index = target_indices[0]
-
-    # Initialize dictionary to store subject level values
-    subject_level_values = {target_column: target_value}
-
-    # Initialize subject level start index
-    subject_level_start = first_target_index
-
-    # Find the last non-NaN row for each subject level
-    for level in range(subject_level_i - 1, 0, -1):  # Loop from subject_level_i-1 to 1
-        column_name = f"Subject Level {level}"
-
-        # Start checking above the previous found index
-        current_index = subject_level_start - 1
-
-        while current_index >= 0 and pd.isna(
-            filtered_df.loc[current_index, column_name]
-        ):
-            current_index -= 1  # Move up while NaN is found
-
-        # Move one row down to get the last valid row in 'Subject Level level'
-        subject_level_start = current_index + 1
-
-        # Ensure we store the correct value at each subject level
-        if subject_level_start in filtered_df.index:
-            subject_level_values[column_name] = filtered_df.loc[
-                subject_level_start - 1, column_name
-            ]
-
-    # Ensure valid slicing range
-    min_start_index = subject_level_start
-
-    if min_start_index < first_target_index:
-        filtered_df = filtered_df.loc[min_start_index:first_target_index]
-    else:
-        filtered_df = pd.DataFrame()
-
-    return subject_level_values, filtered_df
-
-
-def extract_heirarchy(full_code, target_value):
-    # df = pd.read_excel(
-    #     r"C:\Users\mukul.rawat\OneDrive - Candid\Documents\Projects\Gen AI\azure_devops\ask-candid-assistant\PCS_Taxonomy_Definitions_2024.xlsx"
-    # )
-    df = pd.read_excel(r"C:\Users\siqi.deng\Downloads\PCS_Taxonomy_Definitions_2024.xlsx")
-    filtered_df = df[df["PCS Code"].str.startswith(full_code[:2], na=False)]
-    for i in range(1, 5):
-        column_name = f"Subject Level {i}"
-        if (df[column_name].str.strip() == target_value).any():
-            break
-
-    subject_level_values, filtered_df = find_subject_levels(
-        filtered_df, i, target_value
-    )
-    sorted_values = [
-        value
-        for key, value in sorted(
-            subject_level_values.items(), key=lambda x: int(x[0].split()[-1])
-        )
-    ]
-    # Joining values in the required format
-    result = " : ".join(sorted_values)
-    return result
-
-
-class GraphState(TypedDict):
-    query: str = Field(
-        ..., description="The user's query to be processed by the system."
-    )
-    agent_out: str = Field(
-        ...,
-        description="The output generated by the AI agent after processing the query.",
-    )
-    next_step: str = Field(
-        ..., description="The next step in the workflow, determined by query analysis."
-    )
-    es_query: dict = Field(
-        ..., description="The Elasticsearch query generated or used by the agent."
-    )
-
-    es_result: dict = Field(
-        ...,
-        description="The Elasticsearch query result generated or used by the agent.",
-    )
-    pcs_codes: dict = Field(..., description="pcs codes")
-
-
-class AnalysisResult(BaseModel):
-    category: str = Field(..., description="Either 'general' or 'Database'")
-
-
-def agent_factory(llm: LLM) -> AgentExecutor:
-    """
-    Creates and configures an AgentExecutor instance for interacting with Elasticsearch.
-
-    This function initializes an OpenAI GPT-4-based LLM with specific parameters,
-    constructs a prompt tailored for Elasticsearch assistance, and integrates the
-    agent with a set of tools to handle user queries. The agent is designed to work
-    with OpenAI functions for enhanced capabilities.
-
-    Returns:
-        AgentExecutor: Configured agent ready to execute tasks with specified tools,
-                       providing detailed intermediate steps for transparency.
-    """
-
-    # llm = ChatOpenAI(
-    #     model="gpt-4o", temperature=0, api_key=OPENAI["key"], streaming=False
-    # )
-
-    tags_ = []
-    agent = AgentType.OPENAI_FUNCTIONS
-    tags_.append(agent.value if isinstance(agent, AgentType) else agent)
-    # Create the prompt
-    prompt = ChatPromptTemplate.from_messages(
-        [
-            ("system", "You are a helpful elasticsearch assistant"),
-            MessagesPlaceholder(variable_name="chat_history", optional=True),
-            ("human", "{input}"),
-            MessagesPlaceholder(variable_name="agent_scratchpad"),
-        ]
-    )
-
-    # Create the agent
-    agent_obj = create_openai_functions_agent(llm, tools, prompt)
-
-    return AgentExecutor.from_agent_and_tools(
-        agent=agent_obj,
-        tools=tools,
-        tags=tags_,
-        verbose=True,
-        return_intermediate_steps=True,
-    )
-
-
-def agent_factory_claude(llm: LLM) -> AgentExecutor:
-    """
-    Creates and configures an AgentExecutor instance for interacting with Elasticsearch.
-
-    This function initializes an OpenAI GPT-4-based LLM with specific parameters,
-    constructs a prompt tailored for Elasticsearch assistance, and integrates the
-    agent with a set of tools to handle user queries. The agent is designed to work
-    with OpenAI functions for enhanced capabilities.
-
-    Returns:
-        AgentExecutor: Configured agent ready to execute tasks with specified tools,
-                       providing detailed intermediate steps for transparency.
-    """
-
-    # llm = ChatOpenAI(
-    #     model="gpt-4o", temperature=0, api_key=OPENAI["key"], streaming=False
-    # )
-
-    # tags_ = []
-    # agent = AgentType.OPENAI_FUNCTIONS
-    # tags_.append(agent.value if isinstance(agent, AgentType) else agent)
-    # Create the prompt
-    prompt = ChatPromptTemplate.from_messages(
-        [
-            ("system", "You are a helpful elasticsearch assistant"),
-            MessagesPlaceholder(variable_name="chat_history", optional=True),
-            ("human", "{input}"),
-            MessagesPlaceholder(variable_name="agent_scratchpad"),
-        ]
-    )
-
-    agent = create_tool_calling_agent(llm, tools, prompt)
-    agent_executor = AgentExecutor.from_agent_and_tools(
-        agent=agent, tools=tools, verbose=True, return_intermediate_steps=True
-    )
-    # Create the agent
-    return agent_executor
-
-
-# define graph node functions
-def general_query(state: GraphState, llm: LLM) -> GraphState:
-    """
-    Processes a user query using an LLM and updates the graph state with the response.
-
-    Args:
-        state (GraphState): Current graph state containing the user's query.
-        llm (LLM): Language model to process the query.
-
-    Returns:
-        GraphState: Updated state with the LLM's response in "agent_out".
-    """
-    print("> General query")
-    prompt = ChatPromptTemplate.from_template(
-        "Answer based on the user's query: {query}"
-    )
-    chain = prompt | llm
-    response = chain.invoke({"query": state["query"]})
-    if isinstance(response, BaseMessage):
-        state["agent_out"] = response.content
-    else:
-        state["agent_out"] = str(response)
-    return state
-
-
-def database_agent(state: GraphState, llm: LLM) -> GraphState:
-    """
-    Executes a database query using an Elasticsearch agent and updates the graph state.
-
-    The agent queries indices and field names in the Elasticsearch database,
-    selects the appropriate index (`organization_dev_2`), and answers the user's question.
-
-    Args:
-        state (GraphState): Current graph state containing the user's query.
-
-    Returns:
-        GraphState: Updated state with the agent's output in "agent_out" and
-                    the Elasticsearch query in "es_query".
-    """
-
-    print("> database agent")
-    input_data = {
-        "input": f"""
-        You are an Elasticsearch database agent designed to accurately understand and respond to user queries. Follow these steps:
-
-            1. Understand the user query to determine the required information.
-            2. Query the indices in the Elasticsearch database.
-            3. Retrieve the mappings and field names relevant to the query.
-            4. Use the organization_dev_2 index to extract the necessary data.
-            5. Present the response in a clear and natural language format, addressing the user's question directly.
-
-        User's quer:
-        ```{state["query"]}```
-        """
-    }
-    agent_exec = agent_factory_claude(llm)
-    res = agent_exec.invoke(input_data)
-    state["agent_out"] = res["output"]
-
-    es_queries, es_results = {}, {}
-    for i, action in enumerate(res.get("intermediate_steps", []), start=1):
-        if action[0].tool == "elastic_index_search_tool":
-            es_queries[f"query_{i}"] = json.loads(
-                action[0].tool_input.get("query") or "{}"
-            )
-            es_results[f"query_{i}"] = ast.literal_eval(action[-1] or "{}")
-
-    # if len(res["intermediate_steps"]) > 1:
-    #     es_queries = {
-    #         f"query_{i}": action[0].tool_input.get("query", "")
-    #         for i, action in enumerate(res.get("intermediate_steps", []), start=1)
-    #         if action[0].tool == "elastic_index_search_tool"
-    #     }
-
-    #     es_results = {
-    #         f"result_{i}": action[-1]
-    #         for i, action in enumerate(res.get("intermediate_steps", []), start=1)
-    #         if action[0].tool == "elastic_index_search_tool"
-    #     }
-
-    #     state["es_query"] = es_queries
-    #     state["es_result"] = es_results
-    # else:
-    #     state["es_query"] = res["intermediate_steps"][-1][0].tool_input["query"]
-    #     state["es_result"] = {"result": res["intermediate_steps"][-2][-1]}
-
-    state["es_query"] = es_queries
-    state["es_result"] = es_results
-    return state
-
-
-def analyse_query(state: GraphState, llm: LLM) -> GraphState:
-    """
-    Analyzes the user's query to classify it as either general or database-specific
-    and determines the next processing step.
-
-    Args:
-        state (GraphState): Current graph state containing the user's query.
-        llm (LLM): Language model used for query analysis.
-
-    Returns:
-        GraphState: Updated state with the classification result and the
-                    next processing step in "next_step".
-    """
-
-    print("> analyse query")
-    prompt_template = """Your task is to analyze the query  ```{query}``` and classify it in:
-    general: it's a basic general enquiry
-    Database: query which is complicated and would require to go into the database and extract specific information
-    Output format:
-    {{"category": "<your_classification>"}}
-    """
-
-    # Create the prompt
-    prompt = ChatPromptTemplate.from_template(prompt_template)
-
-    # Define the parser
-    parser = PydanticOutputParser(pydantic_object=AnalysisResult)
-
-    # Create the chain
-    chain = RunnableSequence(prompt, llm)
-    # Invoke the chain with the query
-    response = chain.invoke({"query": state["query"]})
-    if "Database" in response.content:
-        state["next_step"] = "es_database_agent"
-    else:
-        state["next_step"] = "general_query"
-    return state
-
-
-def final_answer(state: GraphState, llm: LLM) -> GraphState:
-    """
-    Generates and presents the final response based on the user's query and the AI's output.
-
-    Args:
-        state (GraphState): Current graph state containing the query and AI output.
-        llm (LLM): Language model used to format the final response.
-
-    Returns:
-        GraphState: Updated state with the formatted final answer in "agent_out".
-    """
-
-    print("> Final Answer")
-    prompt_template = """
-    You are a chat agent that takes outputs generated by Elasticsearch and presents them in a conversational, natural language format, as if responding to a user's query.
-
-    Query: ```{query}```
-
-    AI Output: 
-    ```{output}```
-    """
-    prompt = ChatPromptTemplate.from_template(prompt_template)
-    chain = RunnableSequence(prompt, llm)
-    response = chain.invoke({"query": state["query"], "output": state["agent_out"]})
-
-    return {"agent_out": response.content}
-
-
-def build_compute_graph(llm: LLM) -> StateGraph:
-    """
-    Constructs a compute graph for processing user queries using a defined workflow.
-
-    The workflow includes nodes for query analysis, handling general or database-specific queries,
-    and generating the final response. Conditional logic determines the path based on query type.
-
-    Args:
-        llm (LLM): Language model to be used in various nodes for processing queries.
-
-    Returns:
-        StateGraph: Configured compute graph ready for execution.
-    """
-    # Create the workflow
-    workflow = StateGraph(GraphState)
-
-    # Add nodes
-    workflow.add_node("analyse", partial(analyse_query, llm=llm))
-    workflow.add_node("general_query", partial(general_query, llm=llm))
-    workflow.add_node("es_database_agent", partial(database_agent, llm=llm))
-    workflow.add_node("final_answer", partial(final_answer, llm=llm))
-
-    # Set entry point
-    workflow.set_entry_point("analyse")
-
-    # Add conditional edges
-    workflow.add_conditional_edges(
-        "analyse",
-        lambda x: x["next_step"],  # Use the return value of analyse_query directly
-        {"es_database_agent": "es_database_agent", "general_query": "general_query"},
-    )
-
-    # Add edges to end the workflow
-    workflow.add_edge("es_database_agent", "final_answer")
-    workflow.add_edge("general_query", "final_answer")
-    workflow.add_edge("final_answer", END)
-
-    return workflow
-
-
-class ElasticGraph(StateGraph):
-    llm: LLM
-    tools: List[Tool]
-
-    def __init__(self, llm: LLM, tools: List[Tool]):
-        super().__init__(GraphState)
-        self.llm = llm
-        self.tools = tools
-        self.construct_graph()
-
-    def Extract_PCS_Codes(self, state):
-        """Todo: Add Subject heirarchies, Population, Geo"""
-        print("query", state["query"])
-        autocoding_api = AutocodingAPI()
-        autocoding_response = autocoding_api(text=state["query"]).get("data", {})
-        # population_served = autocoding_response.get("population", {})
-        subjects = autocoding_response.get("subject", {})
-        descriptions = []
-        heirarchy_string = []
-        if subjects and isinstance(subjects, list) and "description" in subjects[0]:
-            for subject in subjects:
-                # if subject['description'] in subjects_list:
-                descriptions.append(subject["description"])
-                heirarchy_string.append(
-                    extract_heirarchy(subject["full_code"], subject["description"])
-                )
-        print("descriptions", descriptions)
-
-        populations = autocoding_response.get("population", {})
-        population_dict = []
-        if (
-            populations
-            and isinstance(populations, list)
-            and "description" in populations[0]
-        ):
-            for population in populations:
-                population_dict.append(population["description"])
-        state["pcs_codes"] = {
-            "subject": descriptions,
-            "heirarchy_string": heirarchy_string,
-            "population": population_dict,
-        }
-        print("pcs_codes_new", state["pcs_codes"])
-        return state
-
-    def agent_factory(self) -> AgentExecutor:
-        """
-        Creates and configures an AgentExecutor instance for interacting with Elasticsearch.
-
-        This function initializes an OpenAI GPT-4-based LLM with specific parameters,
-        constructs a prompt tailored for Elasticsearch assistance, and integrates the
-        agent with a set of tools to handle user queries. The agent is designed to work
-        with OpenAI functions for enhanced capabilities.
-
-        Returns:
-            AgentExecutor: Configured agent ready to execute tasks with specified tools,
-                        providing detailed intermediate steps for transparency.
-        """
-
-        # llm = ChatOpenAI(
-        #     model="gpt-4o", temperature=0, api_key=OPENAI["key"], streaming=False
-        # )
-
-        tags_ = []
-        agent = AgentType.OPENAI_FUNCTIONS
-        tags_.append(agent.value if isinstance(agent, AgentType) else agent)
-        # Create the prompt
-        prompt = ChatPromptTemplate.from_messages(
-            [
-                ("system", "You are a helpful elasticsearch assistant"),
-                MessagesPlaceholder(variable_name="chat_history", optional=True),
-                ("human", "{input}"),
-                MessagesPlaceholder(variable_name="agent_scratchpad"),
-            ]
-        )
-
-        # Create the agent
-        agent_obj = create_openai_functions_agent(self.llm, tools, prompt)
-
-        return AgentExecutor.from_agent_and_tools(
-            agent=agent_obj,
-            tools=tools,
-            tags=tags_,
-            verbose=True,
-            return_intermediate_steps=True,
-        )
-
-    def agent_factory_claude(self, pcs_codes, prefix) -> AgentExecutor:
-        """
-        Creates and configures an AgentExecutor instance for interacting with Elasticsearch.
-
-        This function initializes an OpenAI GPT-4-based LLM with specific parameters,
-        constructs a prompt tailored for Elasticsearch assistance, and integrates the
-        agent with a set of tools to handle user queries. The agent is designed to work
-        with OpenAI functions for enhanced capabilities.
-
-        Returns:
-            AgentExecutor: Configured agent ready to execute tasks with specified tools,
-                        providing detailed intermediate steps for transparency.
-        """
-        prompt = ChatPromptTemplate.from_messages(
-            [
-                ("system", f"You are a helpful elasticsearch assistant. {prefix}"),
-                MessagesPlaceholder(variable_name="chat_history", optional=True),
-                ("human", "{input}"),
-                MessagesPlaceholder(variable_name="agent_scratchpad"),
-            ]
-        )
-
-        tools = [
-            # ListIndicesTool(),
-            IndexShowDataTool(),
-            IndexDetailsTool(),
-            create_search_tool(pcs_codes=pcs_codes),
-        ]
-        agent = create_tool_calling_agent(self.llm, tools, prompt)
-
-        agent_executor = AgentExecutor.from_agent_and_tools(
-            agent=agent,
-            tools=tools,
-            verbose=True,
-            return_intermediate_steps=True,
-        )
-        # Create the agent
-        return agent_executor
-
-    def analyse_query(self, state: GraphState) -> GraphState:
-        """
-        Analyzes the user's query to classify it as either general or database-specific
-        and determines the next processing step.
-
-        Args:
-            state (GraphState): Current graph state containing the user's query.
-            llm (LLM): Language model used for query analysis.
-
-        Returns:
-            GraphState: Updated state with the classification result and the
-                        next processing step in "next_step".
-        """
-
-        print("> analyse query")
-        prompt_template = """Your task is to analyze the query  ```{query}``` and classify it in:
-        grant: Grant Index - A query where users seek information about grants, funding opportunities, and grantmakers. This includes inquiries about the purpose of funding, eligibility criteria, application processes, grant recipients, funding amounts, deadlines, and how grants can be used for specific projects or initiatives. Users may also request grants tailored to their unique needs, industries, or social impact goals
-
-        org: Org Index - Query which asks speicific details about the organizations, their mission statement, where they are located
-        Output format:
-        {{"category": "<your_classification>"}}
-        """
-        parser = PydanticOutputParser(pydantic_object=AnalysisResult)
-
-        # Create the prompt
-        prompt = PromptTemplate(
-            template=prompt_template,
-            input_variables=["query"],
-            partial_variables={"format_instructions": parser.get_format_instructions()},
-        )
-        # Create the chain
-        chain = RunnableSequence(prompt, self.llm, parser)
-        # Invoke the chain with the query
-        response = chain.invoke({"query": state["query"]})
-        if response.category == "grant":
-            state["next_step"] = "grant-index"
-        else:
-            state["next_step"] = "org-index"
-        return state
-
-    def grant_index_agent(self, state: GraphState) -> GraphState:
-        print("> Grant Index Agent")
-        # autocoding test
-
-        input_data = {
-            "input": f"""
-            You are an Elasticsearch database agent designed to accurately understand and respond to user queries. Follow these steps:
-
-                1. Understand the user query to determine the required information.
-                2. Query the indices in the Elasticsearch database.
-                3. Retrieve the mappings and field names relevant to the query.
-                4. Use the ``grants_qa_1`` index to extract the necessary data.
-                5. Ensure that you correctly identify the grantmaker (funder) or recipient (funded entity) if mentioned in the query. 
-                   Users may not always provide the exact name, so the Elasticsearch query should accommodate partial or incomplete names 
-                   by searching for relevant keywords.
-                6. Present the response in a clear and natural language format, addressing the user's question directly.
-    
-            Description of some of the fields in the index but rest of the fields which are not here should be easy to understand:
-                *fiscal_year: Year when grantmaker allocates budget for funding and grants. format YYYY 
-                *recipient_state: is abbreviated for eg. NY, FL, CA
-                *recipient_city - Full Name of the City e.g, New York City, Boston 
-                *recipient_country - Country Abbreviation of the recipient organization e.g USA
-
-            Note: Do not include `title`, `program_area`, `text` field in the elastic search query
-            User's query:
-            ```{state["query"]}```
-            """
-        }
-        pcs_codes = state["pcs_codes"]
-        pcs_match_term = ""
-        for pcs_code in pcs_codes["subject"]:
-            if pcs_code != "Philanthropy":
-                pcs_match_term += f"*'pcs_v3.subject.value.name': {pcs_code}* \n"
-
-        for pcs_code in pcs_codes["population"]:
-            if pcs_code != "Other population":
-                pcs_match_term += f"*'pcs_v3.population.value.name': {pcs_code}* \n"
-        print("pcs_match_term", pcs_match_term)
-        prefix = f""" 
-        You are an intelligent agent tasked with generating accurate Elasticsearch DSL queries.
-        Analyze the intent behind the query and determine the appropriate Elasticsearch operations required. 
-        Guidelines for generating right elastic seach query:
-            1. Automatically determine whether to return document hits or aggregation results based on the query structure.
-            2. Use keyword fields instead of text fields for aggregations and sorting to avoid fielddata errors
-            3. Avoid using field.keyword if a keyword field is already present to prevent redundant queries.
-            4. Ensure efficient query execution by selecting appropriate query types for filtering, searching, and aggregating.
-
-         Instruction for pcs_v3 Field-
-            If {pcs_codes['subject']} not empty:
-            Only include all of the following match terms. No other pcs_v3 fields should be added, duplicated, or altered except for those listed below.
-                - {pcs_match_term}
-        """
-        agent_exec = self.agent_factory_claude(
-            pcs_codes=state["pcs_codes"], prefix=prefix
-        )
-        res = agent_exec.invoke(input_data)
-        state["agent_out"] = res["output"]
-        es_queries, es_results = {}, {}
-        for i, action in enumerate(res.get("intermediate_steps", []), start=1):
-            if action[0].tool == "elastic_index_search_tool":
-                print("query", action[0].tool_input.get("query"))
-                es_queries[f"query_{i}"] = json.loads(
-                    action[0].tool_input.get("query") or "{}"
-                )
-                es_results[f"query_{i}"] = ast.literal_eval(action[-1] or "{}")
-
-        state["es_query"] = es_queries
-        state["es_result"] = es_results
-        return state
-
-    def org_index_agent(self, state: GraphState) -> GraphState:
-        """
-        Executes a database query using an Elasticsearch agent and updates the graph state.
-
-        The agent queries indices and field names in the Elasticsearch database,
-        selects the appropriate index (`organization_dev_2`), and answers the user's question.
-
-        Args:
-            state (GraphState): Current graph state containing the user's query.
-
-        Returns:
-            GraphState: Updated state with the agent's output in "agent_out" and
-                        the Elasticsearch query in "es_query".
-        """
-
-        print("> Org Index Agent")
-        mapping_description = """ 
-                    "admin1_code": "state abbreviation"
-                    "admin1_description": "Full name/label of the state"
-                    "city": Full Name of the city with 1st letter being capital for e.g. New York City
-                    "assets": "The assets value of the most recent fiscals available for the organization."
-                    "country_code": "Country abbreviation"
-                    "country_name": "Country name"
-                    "fiscal_year": "The year of the most recent fiscals available for the organization. (YYYY format)"
-                    "mission_statement": "The mission statement of the organization."
-                    "roles": "grantmaker: Indicates the organization gives grants., recipient: Indicates the organization receives grants., company: Indicates the organization is a company/corporation."
-
-            """
-        input_data = {
-            "input": f"""
-            You are an Elasticsearch database agent designed to accurately understand and respond to user queries. Follow these steps:
-
-                1. Understand the user query to determine the required information.
-                2. Query the indices in the Elasticsearch database.
-                3. Retrieve the mappings and field names relevant to the query.
-                4. Use the `organization_qa_ds1` index to extract the necessary data.
-                5. Present the response in a clear and natural language format, addressing the user's question directly.
-            
-                
-            Given Below is mapping description of some of the fields
-            ```{mapping_description}```
-
-
-            User's query:
-            ```{state["query"]}```
-            """
-        }
-
-        pcs_codes = state["pcs_codes"]
-        pcs_match_term = ""
-        for pcs_code in pcs_codes["subject"]:
-            pcs_match_term += f'"taxonomy_descriptions": "{pcs_code}"  \n"'
-
-        print("pcs_match_term", pcs_match_term)
-        prefix = f"""You are an intelligent agent tasked with generating accurate Elasticsearch DSL queries.
-        Analyze the intent behind the query and determine the appropriate Elasticsearch operations required. 
-        Guidelines for generating right elastic seach query:
-            1. Automatically determine whether to return document hits or aggregation results based on the query structure.
-            2. Use keyword fields instead of text fields for aggregations and sorting to avoid fielddata errors
-            3. Avoid using field.keyword if a keyword field is already present to prevent redundant queries.
-            4. Ensure efficient query execution by selecting appropriate query types for filtering, searching, and aggregating.
-            
-            Instructions to use `taxonomy_descriptions` field:
-            If {pcs_codes['subject']} not empty, only add the following match term:
-                    Only add the following `match` term, No other `taxonomy_descriptions` fields should be added, duplicated, or modified except belowIf {pcs_codes['subject']} not empty,
-                    - {pcs_match_term}
-
-
-            Avoid using `ntee_major_description` field in the es query
-
-            """
-        agent_exec = self.agent_factory_claude(
-            pcs_codes=state["pcs_codes"], prefix=prefix
-        )
-        res = agent_exec.invoke(input_data)
-        state["agent_out"] = res["output"]
-
-        es_queries, es_results = {}, {}
-        for i, action in enumerate(res.get("intermediate_steps", []), start=1):
-            if action[0].tool == "elastic_index_search_tool":
-                es_queries[f"query_{i}"] = json.loads(
-                    action[0].tool_input.get("query") or "{}"
-                )
-                es_results[f"query_{i}"] = ast.literal_eval(action[-1] or "{}")
-
-        state["es_query"] = es_queries
-        state["es_result"] = es_results
-        return state
-
-    def final_answer(self, state: GraphState) -> GraphState:
-        """
-        Generates and presents the final response based on the user's query and the AI's output.
-
-        Args:
-            state (GraphState): Current graph state containing the query and AI output.
-            llm (LLM): Language model used to format the final response.
-
-        Returns:
-            GraphState: Updated state with the formatted final answer in "agent_out".
-        """
-
-        print("> Final Answer")
-        prompt_template = """
-        You are a chat agent that takes outputs generated by Elasticsearch and presents them in a conversational, natural language format, as if responding to a user's query.
-
-        Query: ```{query}```
-
-        AI Output: 
-        ```{output}```
-        """
-        prompt = ChatPromptTemplate.from_template(prompt_template)
-        chain = RunnableSequence(prompt, self.llm)
-        response = chain.invoke({"query": state["query"], "output": state["agent_out"]})
-
-        return {"agent_out": response.content}
-
-    def construct_graph(self) -> StateGraph:
-        """
-        Constructs a compute graph for processing user queries using a defined workflow.
-
-        The workflow includes nodes for query analysis, handling general or database-specific queries,
-        and generating the final response. Conditional logic determines the path based on query type.
-
-        Args:
-            llm (LLM): Language model to be used in various nodes for processing queries.
-
-        Returns:
-            StateGraph: Configured compute graph ready for execution.
-        """
-
-        # Add nodes
-        self.add_node("Context_Extraction", self.Extract_PCS_Codes)
-        self.add_node("analyse", self.analyse_query)
-        self.add_node("grant-index", self.grant_index_agent)
-        self.add_node("org-index", self.org_index_agent)
-        self.add_node("final_answer", self.final_answer)
-
-        # Set entry point
-        self.set_entry_point("Context_Extraction")
-        self.add_edge("Context_Extraction", "analyse")
-
-        # Add conditional edges
-        self.add_conditional_edges(
-            "analyse",
-            lambda x: x["next_step"],  # Use the return value of analyse_query directly
-            {"org-index": "org-index", "grant-index": "grant-index"},
-        )
-
-        # Add edges to end the workflow
-        self.add_edge("org-index", "final_answer")
-        self.add_edge("grant-index", "final_answer")
-        self.add_edge("final_answer", END)
-
-
-def build_elastic_graph(llm: LLM, tools: List[Tool]):
-    """Compile Elastic Agent Graph"""
-    elastic_graph = ElasticGraph(llm=llm, tools=tools)
-    graph = elastic_graph.compile()
-    return graph

ask_candid/agents/schema.py

@@ -1,27 +0,0 @@
-from typing import List, Dict, TypedDict, Sequence, Union, Annotated
-
-from langchain_core.messages import BaseMessage
-from langgraph.graph.message import add_messages
-
-
-class Context(TypedDict):
-    """PCS + geonames context payload for common tasks like recommendations.
-    """
-    subject: List[str]
-    population: List[str]
-    geography: List[Union[str, int]]
-
-
-class AgentState(TypedDict):
-    """State of the chat agent for the execution graph(s).
-    """
-    # The add_messages function defines how an update should be processed
-    # Default is to replace. add_messages says "append"
-    messages: Annotated[Sequence[BaseMessage], add_messages]
-    user_input: str
-    org_dict: Dict
-
-    # Recommendation-specific fields
-    intent: str
-    context: Context
-    recommendation: str