Skip to content

Commit 1f3c7cf

Browse files
committed
update sql, neo4j arango examples
1 parent 2cac4bb commit 1f3c7cf

File tree

4 files changed

+208
-5
lines changed

4 files changed

+208
-5
lines changed

examples/data-qa/sql-chat/sql_chat.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,6 +191,7 @@ def main(
191191
show_stats=False,
192192
context_descriptions=context_descriptions, # Add context descriptions to the config
193193
use_schema_tools=schema_tools,
194+
addressing_prefix=SEND_TO,
194195
llm=OpenAIGPTConfig(
195196
chat_model=OpenAIChatModel.GPT4,
196197
),

examples/kg-chat/chat-arangodb.py

Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
"""
2+
Single-agent to use to chat with an existing ArangoDB knowledge-graph (KG) on cloud,
3+
or locally.
4+
If you have an existing ArangoDB instance, you can
5+
chat with it by specifying its URL, username, password, and database name in the dialog.
6+
7+
Run like this (--model is optional, defaults to GPT4o):
8+
9+
python3 examples/kg-chat/chat-arangodb.py --model litellm/claude-3-5-sonnet-20241022
10+
11+
If using litellm, remember to install langroid with the litellm extra, e.g.
12+
pip install "langroid[litellm]"
13+
14+
See these guides for info on setting up langroid to use Open/Local LLMs
15+
and other non-OpenAI LLMs:
16+
- https://langroid.github.io/langroid/tutorials/local-llm-setup/
17+
- https://langroid.github.io/langroid/tutorials/non-openai-llms/
18+
"""
19+
20+
import typer
21+
import os
22+
from typing import Optional
23+
from dotenv import load_dotenv
24+
from rich import print
25+
from rich.console import Console
26+
from rich.prompt import Prompt
27+
28+
import langroid.language_models as lm
29+
from langroid import TaskConfig
30+
from langroid.agent.special.arangodb.arangodb_agent import (
31+
ArangoChatAgentConfig,
32+
ArangoChatAgent,
33+
ArangoSettings,
34+
)
35+
from langroid.utils.constants import SEND_TO
36+
from langroid.agent.chat_document import ChatDocument
37+
from langroid.agent.task import Task
38+
from langroid.utils.configuration import Settings, set_global
39+
from adb_cloud_connector import get_temp_credentials
40+
from arango.client import ArangoClient
41+
from arango_datasets import Datasets
42+
import logging
43+
44+
logging.basicConfig(
45+
level=logging.INFO,
46+
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
47+
force=True, # Add this
48+
)
49+
root_logger = logging.getLogger()
50+
root_logger.setLevel(logging.INFO) # Add this
51+
logger = logging.getLogger(__name__)
52+
53+
console = Console()
54+
app = typer.Typer()
55+
56+
57+
class MyArangoChatAgent(ArangoChatAgent):
58+
def user_response(
59+
self,
60+
msg: Optional[str | ChatDocument] = None,
61+
) -> Optional[ChatDocument]:
62+
response = super().user_response(msg)
63+
if response.content == "r":
64+
65+
self.clear_history(1) # remove all msgs after system msg
66+
n_msgs = len(self.message_history)
67+
assert n_msgs == 1
68+
logger.warning("Reset Agent history, only system msg remains")
69+
# prompt user again
70+
return super().user_response(msg)
71+
72+
return response
73+
74+
75+
@app.command()
76+
def main(
77+
debug: bool = typer.Option(False, "--debug", "-d", help="debug mode"),
78+
model: str = typer.Option("", "--model", "-m", help="model name"),
79+
no_stream: bool = typer.Option(False, "--nostream", "-ns", help="no streaming"),
80+
nocache: bool = typer.Option(False, "--nocache", "-nc", help="don't use cache"),
81+
) -> None:
82+
set_global(
83+
Settings(
84+
debug=debug,
85+
cache=nocache,
86+
stream=not no_stream,
87+
)
88+
)
89+
print(
90+
"""
91+
[blue]Welcome to ArangoDB Knowledge Graph RAG chatbot!
92+
Enter x or q to quit at any point.
93+
"""
94+
)
95+
96+
load_dotenv()
97+
98+
url = Prompt.ask(
99+
"ArangoDB URL (enter 'got' for Game of Thrones dataset) ",
100+
default="https://db.catalog.igvf.org",
101+
)
102+
username = Prompt.ask(
103+
"ArangoDB username ",
104+
default="guest",
105+
)
106+
db = Prompt.ask(
107+
"ArangoDB database ",
108+
default="igvf",
109+
)
110+
pw = Prompt.ask(
111+
"ArangoDB password ",
112+
default="",
113+
)
114+
pw = pw or os.getenv("ARANGODB_PASSWORD")
115+
if url == "got":
116+
print(
117+
"""
118+
No URL supplied, using Game of Thrones dataset from cloud, see here:
119+
https://docs.arangodb.com/3.11/components/tools/arango-datasets/
120+
"""
121+
)
122+
connection = get_temp_credentials(tutorialName="langroid")
123+
client = ArangoClient(hosts=connection["url"])
124+
125+
db = client.db(
126+
connection["dbName"],
127+
connection["username"],
128+
connection["password"],
129+
verify=True,
130+
)
131+
datasets = Datasets(db)
132+
ArangoChatAgent.cleanup_graph_db(db)
133+
assert len(datasets.list_datasets()) > 0, "No datasets found"
134+
135+
DATASET = "GAME_OF_THRONES" # a small dataset
136+
info = datasets.dataset_info(DATASET)
137+
assert info["label"] == DATASET
138+
datasets.load(DATASET, batch_size=100, preserve_existing=False)
139+
arango_settings = ArangoSettings(db=db, client=client)
140+
else:
141+
arango_settings = ArangoSettings(
142+
url=url,
143+
username=username,
144+
database=db,
145+
password=pw,
146+
)
147+
148+
arango_agent = MyArangoChatAgent(
149+
ArangoChatAgentConfig(
150+
name="Arango",
151+
chat_mode=True,
152+
arango_settings=arango_settings,
153+
prepopulate_schema=True,
154+
use_functions_api=False,
155+
use_tools=True,
156+
database_created=True,
157+
llm=lm.OpenAIGPTConfig(
158+
chat_model=model or lm.OpenAIChatModel.GPT4o,
159+
chat_context_length=128_000,
160+
),
161+
human_prompt=(
162+
"Human (respond, or x/q to quit, r to reset history, "
163+
"or hit enter to continue)"
164+
),
165+
)
166+
)
167+
168+
task_config = TaskConfig(addressing_prefix=SEND_TO)
169+
arango_task = Task(
170+
arango_agent,
171+
# user not awaited, UNLESS LLM explicitly addresses user via recipient_tool
172+
interactive=False,
173+
config=task_config,
174+
)
175+
176+
arango_task.run()
177+
178+
# The above runs the app in a continuous chat.
179+
# Alternatively, to set up a task to answer a single query and quit when done:
180+
181+
# set up arango_agent above with chat_mode=False, set up arango_task as above,
182+
# then run the task with a single query, e.g.:
183+
184+
# result = arango_task.run("What is the location of the gene BRCA1?")
185+
186+
# You can have this in a loop with the user, like so:
187+
188+
# while True:
189+
# query = Prompt.ask("Enter your query")
190+
# if query in ["x", "q"]:
191+
# break
192+
# result = arango_task.run(query)
193+
# print(result.content)
194+
195+
196+
if __name__ == "__main__":
197+
app()

examples/kg-chat/chat-neo4j.py

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,10 @@
44
If you have an existing Neo4j db on Aura (or possibly elsewhere, e.g. locally), you can
55
chat with it by specifying its URI, username, password, and database name in the dialog.
66
7-
If you don't have an existing Neo4j db, and want to try this script, you can populate
7+
You can chose the defaults in the dialog, in which case it will use the
8+
freely available Movies database.
9+
10+
Or, you can populate
811
an empty Neo4j db with the cypher queries in the file `movies.cypher` in this folder.
912
1013
See info on getting setup with Neo4j here:
@@ -64,19 +67,19 @@ def main(
6467

6568
uri = Prompt.ask(
6669
"Neo4j URI ",
67-
default="neo4j+s://8d105d28.databases.neo4j.io",
70+
default="neo4j+s://demo.neo4jlabs.com",
6871
)
6972
username = Prompt.ask(
7073
"No4j username ",
71-
default="neo4j",
74+
default="movies",
7275
)
7376
db = Prompt.ask(
7477
"Neo4j database ",
75-
default="neo4j",
78+
default="movies",
7679
)
7780
pw = Prompt.ask(
7881
"Neo4j password ",
79-
default="",
82+
default="movies",
8083
)
8184
pw = pw or os.getenv("NEO4J_PASSWORD")
8285
neo4j_settings = Neo4jSettings(uri=uri, username=username, database=db, password=pw)
@@ -88,6 +91,7 @@ def main(
8891
neo4j_config = Neo4jChatAgentConfig(
8992
neo4j_settings=neo4j_settings,
9093
llm=llm_config,
94+
chat_mode=True,
9195
)
9296

9397
neo4j_agent = Neo4jChatAgent(neo4j_config)

examples/kg-chat/dependency_chatbot.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -232,6 +232,7 @@ def main(
232232

233233
dependency_agent = DependencyGraphAgent(
234234
config=Neo4jChatAgentConfig(
235+
chat_mode=True,
235236
neo4j_settings=neo4j_settings,
236237
show_stats=False,
237238
use_tools=tools,

0 commit comments

Comments
 (0)