SQL 도구 LangGraph 샘플 코드가 있는 ReAct 에이전트

이 LangGraph 샘플 코드를 사용하여 ReAct 에이전트에서 SQL 툴을 테스트하고 디버그할 수 있습니다.

from aidputils.agents.toolkit.tool_helper import create_langgraph_tool
from aidputils.agents.toolkit.agent_helper import init_oci_llm, pre_invoke_setup
from aidputils.agents.toolkit.configs import AIDPToolConf, OCIAIConf, ModelArgs
from langgraph.prebuilt import create_react_agent
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
import logging
import json
import os

compartment_id = os.getenv('AIDP_USER_COMPARTMENT_ID')
compartment_id="<your-compartment-ocid>"
logger = logging.getLogger('agent_with_sql_tool')
checkpointer = globals().get("checkpointer", None)

########## Guardrails Configuration ################
guardrails_config = {
    "name" : "Default Guardrails",
    "description" : "Default empty guardrails configuration",
    "policies" : [ ]
  }
########## End Guardrails Configuration ############

#####
##### Start SQL Tool configuration

sql_params = [ {
  "name" : "MAX_SALARY",
  "type" : "string",
  "description" : "Maximum salary",
  "defaultValue" : "50000"
} ]


conf = {
    "catalogKey": "adw23ai_phx",
    "schemaKey": "gold",
    "query": """ Select * from (
  Select 101 employee_id, 'John' first_name, 'Doe' last_name, '[e-mail]' email_address, 75000 salary from DUAL
  UNION
  Select 102 employee_id, 'Jane' first_name, 'Smith' last_name, '[e-mail]' email_address, 100000 salary from DUAL
  UNION
  Select 103 employee_id, 'Peter' first_name, 'Jones' last_name, '[e-mail]' email_address, 45000 salary from DUAL
) employees where salary >= {{MAX_SALARY}}
"""
}

sql_conf= AIDPToolConf(name="query_employees",
                                  description= "Query employees using SQL Tool ",
                                  tool_class = "SQLTool", conf=conf, params=sql_params)
sql_tool = create_langgraph_tool(sql_conf.model_dump())


##### End SQL Tool configuration

##### Start tool List#############
tools_agent1 = [sql_tool]
##### End tool List#############


model_args = {}

llm_conf = OCIAIConf(model_provider='generic',
                     compartment_id='<your-compartment-ocid>',
                     model_args=model_args,
                     endpoint='https://inference.generativeai.<oci-region>.oci.oraclecloud.com',
                     model_id='xai.grok-4',
                     guardrails_config=guardrails_config)

## Agent class definition
class AgentWithSQLTool:
  def __init__(self) -> None:
    self.agent = None

  def setup(self) -> None:
    logger.info(llm_conf)
    oci_llm = init_oci_llm(llm_conf)
    system_prompt = """
Be a helpful assistant.
"""

    try:
      if checkpointer:
        self.agent = create_react_agent(model=oci_llm, tools=tools_agent1, prompt=system_prompt, debug=True, checkpointer= checkpointer)
      else:
        self.agent  = create_react_agent(model=oci_llm, tools=tools_agent1, prompt=system_prompt, debug=True)
    except Exception as e:
      # Fallback compile without checkpointer if wiring fails
      self.agent = create_react_agent(model=oci_llm, tools=tools_agent1, prompt=system_prompt, debug=True)
      logger.warning(f"Checkpointer could not be initialized {e}")
    logger.info(f"Setup for agent completed {self.agent}")

  async def invoke(self, user_query: str, **kwargs):
    config = pre_invoke_setup(**kwargs)
    user_message = HumanMessage(content=user_query)
    message = {"messages": [dict(user_message)]}
    try:
      return await self.agent.ainvoke(input=message, config = config)
    except Exception as e:
      import traceback
      logger.error(f"Exception while calling invoke {e}", exc_info=True)
      print("Stack trace:\n", traceback.format_exc())




import asyncio
async def main():
    # Instantiate and initialize the agent
    test_agent = AgentWithSQLTool()
    test_agent.setup()

    # You can customize this user query or prompt for input
    user_query = "Which employees have a salary greater than 50000"

    # Run the asynchronous invoke method and print the result
    result = await test_agent.invoke(user_query)
    print("Agent response:", result)

if __name__ == "__main__":
    asyncio.run(main())