-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmocks.py
43 lines (33 loc) · 1.5 KB
/
mocks.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
from langchain.agents.react.base import DocstoreExplorer
from langchain.llms.base import BaseLLM
def reactLLMMock(prompt: str) -> str:
last_line = prompt.split('\n')[-1].strip()
last_action = last_line.split(' ')[0].lower()
if last_action == 'thought':
return 'It does not mention the eastern sector. So I need to look up eastern sector.'
elif last_action == 'action':
return 'Lookup[eastern sector]'
else:
raise Exception('Invalid action type')
def reflectLLMMock(prompt: str) -> str:
return "Last time i should have answered correctly"
class LLMMock(BaseLLM):
def __init__(self):
...
def __call__(self, prompt: str) -> str:
if prompt.split('\n')[0].split(' ')[0] == 'Solve':
return reactLLMMock(prompt)
elif prompt.split('\n')[0].split(' ')[0] == 'You':
return reflectLLMMock(prompt)
else:
raise Exception("Invalid LLM prompt")
def get_num_tokens(self, text: str) -> int:
return 0
class DocStoreExplorerMock(DocstoreExplorer):
def __init__(self):
self.summary = "The Colorado orogeny was an episode of mountain building (an orogeny) in Colorado and surrounding areas."
self.body = "(Result 1 / 1) The eastern sector extends into the High Plains and is called the Central Plains orogeny."
def search(self, search: str, sents: int = 5) -> str:
return self.summary
def lookup(self, term: str) -> str:
return self.body