-
-
Notifications
You must be signed in to change notification settings - Fork 220
/
oai-retrieval-assistant.py
73 lines (57 loc) · 1.93 KB
/
oai-retrieval-assistant.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
"""
Use OpenAI Assistant with Retrieval tool + file to answer questions.
Run like this:
python3 examples/docqa/oai-retrieval-assistant.py
"""
import typer
from rich import print
from rich.prompt import Prompt
import os
import tempfile
from langroid.agent.openai_assistant import (
OpenAIAssistantConfig,
OpenAIAssistant,
AssistantTool,
)
from langroid.parsing.url_loader import URLLoader
from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
from langroid.agent.task import Task
from langroid.utils.logging import setup_colored_logging
app = typer.Typer()
setup_colored_logging()
os.environ["TOKENIZERS_PARALLELISM"] = "false"
@app.command()
def chat() -> None:
reuse = (
Prompt.ask(
"Reuse existing assistant, threads if available? (y/n)",
default="y",
)
== "y"
)
cfg = OpenAIAssistantConfig(
llm=OpenAIGPTConfig(chat_model=OpenAIChatModel.GPT4o),
use_cached_assistant=reuse,
use_cached_thread=reuse,
system_message="Answer questions based on the provided document.",
)
agent = OpenAIAssistant(cfg)
print("[blue]Welcome to the retrieval chatbot!")
path = Prompt.ask("Enter a URL or file path")
# if path is a url, use UrlLoader to get text as a document
if path.startswith("http"):
text = URLLoader([path]).load()[0].content
# save text to a temp file
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
f.write(text)
f.close()
# get the filename
path = f.name
agent.add_assistant_tools([AssistantTool(type="retrieval")])
if path: # may be empty if continuing from previous session
agent.add_assistant_files([path])
print("[cyan]Enter x or q to quit")
task = Task(agent)
task.run("Please help me with questions about the document I provided")
if __name__ == "__main__":
app()