DB-GPT/examples/agents/sql_agent_dialogue_example.py
明天 0541d1494c
feat(datasource):add oceanbase support (#1622)
Co-authored-by: csunny <cfqsunny@163.com>
Co-authored-by: aries_ckt <916701291@qq.com>
2024-06-13 15:13:50 +08:00

64 lines
1.7 KiB
Python

"""Agents: single agents about CodeAssistantAgent?
Examples:
Execute the following command in the terminal:
Set env params.
.. code-block:: shell
export OPENAI_API_KEY=sk-xx
export OPENAI_API_BASE=https://xx:80/v1
run example.
..code-block:: shell
python examples/agents/single_agent_dialogue_example.py
"""
import asyncio
import os
from dbgpt.agent import AgentContext, AgentMemory, LLMConfig, UserProxyAgent
from dbgpt.agent.expand.data_scientist_agent import DataScientistAgent
from dbgpt.agent.resource import SQLiteDBResource
from dbgpt.util.tracer import initialize_tracer
current_dir = os.getcwd()
parent_dir = os.path.dirname(current_dir)
test_plugin_dir = os.path.join(parent_dir, "test_files")
initialize_tracer("/tmp/agent_trace.jsonl", create_system_app=True)
async def main():
from dbgpt.model.proxy import OpenAILLMClient
llm_client = OpenAILLMClient(model_alias="gpt-3.5-turbo")
context: AgentContext = AgentContext(conv_id="test456")
agent_memory = AgentMemory()
sqlite_resource = SQLiteDBResource("SQLite Database", f"{test_plugin_dir}/dbgpt.db")
user_proxy = await UserProxyAgent().bind(agent_memory).bind(context).build()
sql_boy = (
await DataScientistAgent()
.bind(context)
.bind(LLMConfig(llm_client=llm_client))
.bind(sqlite_resource)
.bind(agent_memory)
.build()
)
await user_proxy.initiate_chat(
recipient=sql_boy,
reviewer=user_proxy,
message="当前库有那些表",
)
## dbgpt-vis message infos
print(await agent_memory.gpts_memory.one_chat_completions("test456"))
if __name__ == "__main__":
asyncio.run(main())