|
| 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() |
0 commit comments