-
-
Notifications
You must be signed in to change notification settings - Fork 209
/
chat-search.py
219 lines (181 loc) · 7.06 KB
/
chat-search.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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
"""
Single-agent question-answering system that has access to DuckDuckGo (DDG) Search when needed,
and in case a DDG Search is used, ingests contents into a vector-db,
and uses Retrieval Augmentation to answer the question.
Run like this:
python3 examples/docqa/chat-search.py
Optional args:
-nc : turn off caching (i.e. don't retrieve cached LLM responses)
-d: debug mode, to show all intermediate results
-f: use OpenAI functions api instead of tools
-m <model_name>: (e.g. -m ollama/mistral:7b-instruct-v0.2-q4_K_M)
(defaults to GPT4-Turbo if blank)
(See here for guide to using local LLMs with Langroid:)
https://langroid.github.io/langroid/tutorials/local-llm-setup/
"""
import re
from typing import List
from rich import print
from rich.prompt import Prompt
import langroid.language_models as lm
from langroid.agent.tool_message import ToolMessage
from langroid.agent.chat_agent import ChatAgent, ChatDocument
from langroid.agent.special.doc_chat_agent import (
DocChatAgent,
DocChatAgentConfig,
)
from langroid.parsing.web_search import duckduckgo_search
from langroid.agent.task import Task
from langroid.utils.constants import NO_ANSWER
from langroid.utils.configuration import set_global, Settings
from fire import Fire
class RelevantExtractsTool(ToolMessage):
request = "relevant_extracts"
purpose = "Get docs/extracts relevant to the <query>"
query: str
@classmethod
def examples(cls) -> List["ToolMessage"]:
return [
cls(query="when was the Mistral LLM released?"),
]
@classmethod
def instructions(cls) -> str:
return """
IMPORTANT: You must include an ACTUAL query in the `query` field,
"""
class RelevantSearchExtractsTool(ToolMessage):
request = "relevant_search_extracts"
purpose = "Get docs/extracts relevant to the <query> from a web search"
query: str
num_results: int = 3
@classmethod
def examples(cls) -> List["ToolMessage"]:
return [
cls(
query="when was the Mistral LLM released?",
num_results=3,
),
]
@classmethod
def instructions(cls) -> str:
return """
IMPORTANT: You must include an ACTUAL query in the `query` field,
"""
class DDGSearchDocChatAgent(DocChatAgent):
tried_vecdb: bool = False
def llm_response(
self,
query: None | str | ChatDocument = None,
) -> ChatDocument | None:
return ChatAgent.llm_response(self, query)
def relevant_extracts(self, msg: RelevantExtractsTool) -> str:
"""Get docs/extracts relevant to the query, from vecdb"""
self.tried_vecdb = True
query = msg.query
_, extracts = self.get_relevant_extracts(query)
if len(extracts) == 0:
return """
No extracts found! You can try doing a web search with the
`relevant_search_extracts` tool/function-call.
"""
return "\n".join(str(e) for e in extracts)
def relevant_search_extracts(self, msg: RelevantSearchExtractsTool) -> str:
"""Get docs/extracts relevant to the query, from a web search"""
if not self.tried_vecdb and len(self.original_docs) > 0:
return "Please try the `relevant_extracts` tool, before using this tool"
self.tried_vecdb = False
query = msg.query
num_results = msg.num_results
results = duckduckgo_search(query, num_results)
links = [r.link for r in results]
self.config.doc_paths = links
self.ingest()
_, extracts = self.get_relevant_extracts(query)
return "\n".join(str(e) for e in extracts)
def main(
debug: bool = False,
nocache: bool = False,
model: str = "",
fn_api: bool = True,
) -> None:
set_global(
Settings(
debug=debug,
cache=not nocache,
)
)
print(
"""
[blue]Welcome to the Internet Search chatbot!
I will try to answer your questions, relying on (full content of links from)
Duckduckgo (DDG) Search when needed.
Enter x or q to quit, or ? for evidence
"""
)
system_msg = Prompt.ask(
"""
[blue] Tell me who I am (give me a role) by completing this sentence:
You are...
[or hit enter for default]
[blue] Human
""",
default="a helpful assistant.",
)
system_msg = re.sub("you are", "", system_msg, flags=re.IGNORECASE)
llm_config = lm.OpenAIGPTConfig(
chat_model=model or lm.OpenAIChatModel.GPT4o,
# or, other possibilities for example:
# "litellm/bedrock/anthropic.claude-instant-v1"
# "ollama/llama2"
# "local/localhost:8000/v1"
# "local/localhost:8000"
chat_context_length=2048, # adjust based on model
)
config = DocChatAgentConfig(
use_functions_api=fn_api,
use_tools=not fn_api,
llm=llm_config,
system_message=f"""
{system_msg} You will try your best to answer my questions,
in this order of preference:
1. If you can answer from your own knowledge, simply return the answer
2. Otherwise, ask me for some relevant text, and I will send you. Use the
`relevant_extracts` tool/function-call for this purpose. Once you receive
the text, you can use it to answer my question.
If I say {NO_ANSWER}, it means I found no relevant docs, and you can try
the next step, using a web search.
3. If you are still unable to answer, you can use the `relevant_search_extracts`
tool/function-call to get some text from a web search. Once you receive the
text, you can use it to answer my question.
4. If you still can't answer, simply say {NO_ANSWER}
Remember to always FIRST try `relevant_extracts` to see if there are already
any relevant docs, before trying web-search with `relevant_search_extracts`.
Be very concise in your responses, use no more than 1-2 sentences.
When you answer based on provided documents, be sure to show me
the SOURCE(s) and EXTRACT(s), for example:
SOURCE: https://www.wikihow.com/Be-a-Good-Assistant-Manager
EXTRACT: Be a Good Assistant ... requires good leadership skills.
For the EXTRACT, ONLY show up to first 3 words, and last 3 words.
""",
)
agent = DDGSearchDocChatAgent(config)
agent.enable_message(RelevantExtractsTool)
agent.enable_message(RelevantSearchExtractsTool)
collection_name = Prompt.ask(
"Name a collection to use",
default="docqa-chat-search",
)
replace = (
Prompt.ask(
"Would you like to replace (i.e. erase) this collection?",
choices=["y", "n"],
default="n",
)
== "y"
)
print(f"[red]Using {collection_name}")
agent.vecdb.set_collection(collection_name, replace=replace)
task = Task(agent)
task.run("Can you help me answer some questions, possibly using web search?")
if __name__ == "__main__":
Fire(main)