ahuang11 commited on
Commit
af19f4e
1 Parent(s): 83b2eea

Delete utils

Browse files
Files changed (4) hide show
  1. utils/__init__.py +0 -0
  2. utils/ai.py +0 -290
  3. utils/stream.py +0 -20
  4. utils/utils.py +0 -70
utils/__init__.py DELETED
File without changes
utils/ai.py DELETED
@@ -1,290 +0,0 @@
1
- # pylint: disable=W0707
2
- # pylint: disable=W0719
3
-
4
- import os
5
- import json
6
- import tiktoken
7
- import openai
8
- from openai import OpenAI
9
- import requests
10
-
11
- from constants.cli import OPENAI_MODELS
12
- from constants.ai import SYSTEM_PROMPT, PROMPT, API_URL
13
-
14
-
15
- def retrieve(query, k=10, filters=None):
16
- """Retrieves and returns dict.
17
-
18
- Args:
19
- query (str): User query to pass in
20
- k (int, optional): number of results passed back. Defaults to 10.
21
- filters (dict, optional): Filters to apply to the query. You can filter based off
22
- any piece of metadata by passing in a dict of the format {metadata_name: filter_value}
23
- ie {"library_id": "1234"}.
24
-
25
- See the README for more details:
26
- https://github.com/fleet-ai/context/tree/main#using-fleet-contexts-rich-metadata
27
-
28
- Returns:
29
- list: List of queried results
30
- """
31
-
32
- url = f"{API_URL}/query"
33
- params = {
34
- "query": query,
35
- "dataset": "python_libraries",
36
- "n_results": k,
37
- "filters": filters,
38
- }
39
- return requests.post(url, json=params, timeout=120).json()
40
-
41
-
42
- def retrieve_context(query, k=10, filters=None):
43
- """Gets the context from our libraries vector db for a given query.
44
-
45
- Args:
46
- query (str): User input query
47
- k (int, optional): number of retrieved results. Defaults to 10.
48
- """
49
-
50
- # First, we query the API
51
- responses = retrieve(query, k=k, filters=filters)
52
-
53
- # Then, we build the prompt_with_context string
54
- prompt_with_context = ""
55
- for response in responses:
56
- prompt_with_context += f"\n\n### Context {response['metadata']['url']} ###\n{response['metadata']['text']}"
57
- return {"role": "user", "content": prompt_with_context}
58
-
59
-
60
- def construct_prompt(
61
- messages,
62
- context_message,
63
- model="gpt-4-1106-preview",
64
- cite_sources=True,
65
- context_window=3000,
66
- ):
67
- """
68
- Constructs a RAG (Retrieval-Augmented Generation) prompt by balancing the token count of messages and context_message.
69
- If the total token count exceeds the maximum limit, it adjusts the token count of each to maintain a 1:1 proportion.
70
- It then combines both lists and returns the result.
71
-
72
- Parameters:
73
- messages (List[dict]): List of messages to be included in the prompt.
74
- context_message (dict): Context message to be included in the prompt.
75
- model (str): The model to be used for encoding, default is "gpt-4-1106-preview".
76
-
77
- Returns:
78
- List[dict]: The constructed RAG prompt.
79
- """
80
- # Get the encoding; default to cl100k_base
81
- if model in OPENAI_MODELS:
82
- encoding = tiktoken.encoding_for_model(model)
83
- else:
84
- encoding = tiktoken.get_encoding("cl100k_base")
85
-
86
- # 1) calculate tokens
87
- reserved_space = 1000
88
- max_messages_count = int((context_window - reserved_space) / 2)
89
- max_context_count = int((context_window - reserved_space) / 2)
90
-
91
- # 2) construct prompt
92
- prompts = messages.copy()
93
- prompts.insert(0, {"role": "system", "content": SYSTEM_PROMPT})
94
- if cite_sources:
95
- prompts.insert(-1, {"role": "user", "content": PROMPT})
96
-
97
- # 3) find how many tokens each list has
98
- messages_token_count = len(
99
- encoding.encode(
100
- "\n".join(
101
- [
102
- f"<|im_start|>{message['role']}\n{message['content']}<|im_end|>"
103
- for message in prompts
104
- ]
105
- )
106
- )
107
- )
108
- context_token_count = len(
109
- encoding.encode(
110
- f"<|im_start|>{context_message['role']}\n{context_message['content']}<|im_end|>"
111
- )
112
- )
113
-
114
- # 4) Balance the token count for each
115
- if (messages_token_count + context_token_count) > (context_window - reserved_space):
116
- # context has more than limit, messages has less than limit
117
- if (messages_token_count < max_messages_count) and (
118
- context_token_count > max_context_count
119
- ):
120
- max_context_count += max_messages_count - messages_token_count
121
- # messages has more than limit, context has less than limit
122
- elif (messages_token_count > max_messages_count) and (
123
- context_token_count < max_context_count
124
- ):
125
- max_messages_count += max_context_count - context_token_count
126
-
127
- # 5) Cut each list to the max count
128
-
129
- # Cut down messages
130
- while messages_token_count > max_messages_count:
131
- removed_encoding = encoding.encode(
132
- f"<|im_start|>{prompts[1]['role']}\n{prompts[1]['content']}<|im_end|>"
133
- )
134
- messages_token_count -= len(removed_encoding)
135
- if messages_token_count < max_messages_count:
136
- prompts = (
137
- [prompts[0]]
138
- + [
139
- {
140
- "role": prompts[1]["role"],
141
- "content": encoding.decode(
142
- removed_encoding[
143
- : min(
144
- int(max_messages_count -
145
- messages_token_count),
146
- len(removed_encoding),
147
- )
148
- ]
149
- )
150
- .replace("<|im_start|>", "")
151
- .replace("<|im_end|>", ""),
152
- }
153
- ]
154
- + prompts[2:]
155
- )
156
- else:
157
- prompts = [prompts[0]] + prompts[2:]
158
-
159
- # Cut down context
160
- if context_token_count > max_context_count:
161
- # Taking a proportion of the content chars length
162
- reduced_chars_length = int(
163
- len(context_message["content"]) *
164
- (max_context_count / context_token_count)
165
- )
166
- context_message["content"] = context_message["content"][:reduced_chars_length]
167
-
168
- # 6) Combine both lists
169
- prompts.insert(-1, context_message)
170
-
171
- return prompts
172
-
173
-
174
- def get_remote_chat_response(messages, model="gpt-4-1106-preview"):
175
- """
176
- Returns a streamed OpenAI chat response.
177
-
178
- Parameters:
179
- messages (List[dict]): List of messages to be included in the prompt.
180
- model (str): The model to be used for encoding, default is "gpt-4-1106-preview".
181
-
182
- Returns:
183
- str: The streamed OpenAI chat response.
184
- """
185
- client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
186
-
187
- try:
188
- response = client.chat.completions.create(
189
- model=model, messages=messages, temperature=0.2, stream=True
190
- )
191
-
192
- for chunk in response:
193
- current_context = chunk.choices[0].delta.content
194
- yield current_context
195
-
196
- except openai.AuthenticationError as error:
197
- print("401 Authentication Error:", error)
198
- raise Exception(
199
- "Invalid OPENAI_API_KEY. Please re-run with a valid key.")
200
-
201
- except Exception as error:
202
- print("Streaming Error:", error)
203
- raise Exception("Internal Server Error")
204
-
205
-
206
- def get_other_chat_response(messages, model="local-model"):
207
- """
208
- Returns a streamed chat response from a local server.
209
-
210
- Parameters:
211
- messages (List[dict]): List of messages to be included in the prompt.
212
- model (str): The model to be used for encoding, default is "gpt-4-1106-preview".
213
-
214
- Returns:
215
- str: The streamed chat response.
216
- """
217
- try:
218
- if model == "local-model":
219
- url = "http://localhost:1234/v1/chat/completions"
220
- headers = {"Content-Type": "application/json"}
221
- data = {
222
- "messages": messages,
223
- "temperature": 0.2,
224
- "max_tokens": -1,
225
- "stream": True,
226
- }
227
- response = requests.post(
228
- url, headers=headers, data=json.dumps(data), stream=True, timeout=120
229
- )
230
-
231
- if response.status_code == 200:
232
- for chunk in response.iter_content(chunk_size=None):
233
- decoded_chunk = chunk.decode()
234
- if (
235
- "data:" in decoded_chunk
236
- and decoded_chunk.split("data:")[1].strip()
237
- ): # Check if the chunk is not empty
238
- try:
239
- chunk_dict = json.loads(
240
- decoded_chunk.split("data:")[1].strip()
241
- )
242
- yield chunk_dict["choices"][0]["delta"].get("content", "")
243
- except json.JSONDecodeError:
244
- pass
245
- else:
246
- print(f"Error: {response.status_code}, {response.text}")
247
- raise Exception("Internal Server Error")
248
- else:
249
- if not os.environ.get("OPENROUTER_API_KEY"):
250
- raise Exception(
251
- f"For non-OpenAI models, like {model}, set your OPENROUTER_API_KEY."
252
- )
253
-
254
- response = requests.post(
255
- url="https://openrouter.ai/api/v1/chat/completions",
256
- headers={
257
- "Authorization": f"Bearer {os.environ.get('OPENROUTER_API_KEY')}",
258
- "HTTP-Referer": os.environ.get(
259
- "OPENROUTER_APP_URL", "https://fleet.so/context"
260
- ),
261
- "X-Title": os.environ.get("OPENROUTER_APP_TITLE", "Fleet Context"),
262
- "Content-Type": "application/json",
263
- },
264
- data=json.dumps(
265
- {"model": model, "messages": messages, "stream": True}),
266
- stream=True,
267
- timeout=120,
268
- )
269
- if response.status_code == 200:
270
- for chunk in response.iter_lines():
271
- decoded_chunk = chunk.decode("utf-8")
272
- if (
273
- "data:" in decoded_chunk
274
- and decoded_chunk.split("data:")[1].strip()
275
- ): # Check if the chunk is not empty
276
- try:
277
- chunk_dict = json.loads(
278
- decoded_chunk.split("data:")[1].strip()
279
- )
280
- yield chunk_dict["choices"][0]["delta"].get("content", "")
281
- except json.JSONDecodeError:
282
- pass
283
- else:
284
- print(f"Error: {response.status_code}, {response.text}")
285
- raise Exception("Internal Server Error")
286
-
287
- except requests.exceptions.RequestException as error:
288
- print("Request Error:", error)
289
- raise Exception(
290
- "Invalid request. Please check your request parameters.")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
utils/stream.py DELETED
@@ -1,20 +0,0 @@
1
- from rich.box import MINIMAL
2
- from rich.live import Live
3
- from rich.panel import Panel
4
- from rich.console import Console
5
- from rich.markdown import Markdown
6
-
7
-
8
- class TextStream:
9
- def __init__(self):
10
- self.live = Live(console=Console(), auto_refresh=False)
11
- self.live.start()
12
-
13
- def print_stream(self, message):
14
- markdown = Markdown(message.strip() + "●")
15
- panel = Panel(markdown, box=MINIMAL)
16
- self.live.update(panel)
17
- self.live.refresh()
18
-
19
- def end_stream(self):
20
- self.live.stop()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
utils/utils.py DELETED
@@ -1,70 +0,0 @@
1
- import re
2
- import traceback
3
-
4
- from rich import print as rprint
5
- from rich.console import Console
6
- from rich.markdown import Markdown
7
- from rich.text import Text
8
- from rich.rule import Rule
9
- from rich.table import Table
10
- from rich.panel import Panel
11
-
12
-
13
- console = Console()
14
-
15
-
16
- def print_markdown(message):
17
- for line in message.split("\n"):
18
- line = line.strip()
19
- if line == "":
20
- print("")
21
- elif line == "---":
22
- rprint(Rule(style="white"))
23
- elif line.startswith("!!!"):
24
- rprint(Text(line[3:], style="#D5D7FB"))
25
- else:
26
- rprint(Markdown(line))
27
-
28
- if "\n" not in message and message.startswith(">"):
29
- print("")
30
-
31
-
32
- def print_exception(exc_type, exc_value, traceback_obj):
33
- traceback_details = traceback.extract_tb(traceback_obj)
34
- for filename, lineno, funcname, text in traceback_details:
35
- console.print(
36
- f"File: {filename}, Line: {lineno}, Func: {funcname}, Text: {text}"
37
- )
38
- console.print(f"{exc_type.__name__}: {exc_value}")
39
-
40
-
41
- def extract_code_blocks(message):
42
- pattern = r"```python\n(.*?)```"
43
- matches = re.findall(pattern, message, re.DOTALL)
44
- return "\n".join(matches)
45
-
46
-
47
- def print_help():
48
- table = Table(show_header=True, header_style="bold magenta")
49
- table.add_column("Command")
50
- table.add_column("Description")
51
-
52
- # Add rows to the table for each command
53
- table.add_row("-k, --k_value", "Number of chunks to return")
54
- table.add_row(
55
- "-l, --libraries",
56
- "Limit your chat to a list of libraries. Usage: -l library1 library2 library3",
57
- )
58
- table.add_row(
59
- "-m, --model", "Specify the model. Default: gpt-4-1106-preview (gpt-4-turbo)"
60
- )
61
- table.add_row(
62
- "-c, --cite_sources", "Determines whether or not the AI model cites its sources"
63
- )
64
- table.add_row("-h, --help", "Help")
65
-
66
- # Create a panel with the table
67
- panel = Panel(table, title="Help", border_style="blue")
68
-
69
- # Print the panel
70
- rprint(panel)