cameron.stream
[READ-ONLY] Mirror of https://github.com/just-cameron/mongo-hack.
9.8 kB
336 lines
1from pathlib import Path
2from fastapi import Depends, HTTPException, status
3from fastapi.responses import FileResponse
4from modal import Stub, web_endpoint, Image, Secret, method, NetworkFileSystem
5import os
6from typing import Dict
7import json
8
9from starlette.requests import Request
10
11APP_NAME = "mongo-hack"
12
13stub = Stub(name=APP_NAME)
14RESULTS_DIR = "/results"
15results_volume = NetworkFileSystem.from_name(f"{APP_NAME}-results-vol", create_if_missing=True)
16
17image = (
18 Image.debian_slim(
19 python_version="3.10"
20 ).pip_install(
21 "replicate",
22 "python-jose",
23 "requests",
24 "openai",
25 "pymongo",
26 "scikit-learn",
27 "numpy"
28 )
29)
30
31@stub.function(image=image, keep_warm=1, secrets=[Secret.from_name("mongo-hack-secret")])
32@web_endpoint(method="GET")
33def fetch_doc_text(
34 docId: str,
35 payload: Dict,
36 request: Request
37):
38
39 print(docId)
40
41 from pymongo import MongoClient
42 from bson.objectid import ObjectId
43
44
45 # Connect to MongoDB Atlas
46 client = MongoClient("mongodb+srv://cameron:testing123@mongo-hackathon.4lo2sdu.mongodb.net/")
47 db = client["mongo-hackathon"]
48 collection = db["documents"]
49
50 print(collection)
51
52 # Specify the document ID
53 document_id = docId
54
55 # Find the document by its ID
56 document = collection.find_one({'_id': ObjectId(document_id)})
57
58 if document:
59 print('Document found:')
60 return document['text']
61 else:
62 print('Document not found.')
63 return document['Document not found.']
64
65
66@stub.function(image=image, keep_warm=1, secrets=[Secret.from_name("mongo-hack-secret")])
67@web_endpoint(method="POST")
68def handle_request(
69 payload: Dict,
70 request: Request
71):
72
73 from pymongo import MongoClient
74 client = MongoClient("mongodb+srv://cameron:testing123@mongo-hackathon.4lo2sdu.mongodb.net/")
75 db = client["mongo-hackathon"]
76 print(db)
77
78 collection = db["documents"]
79 input_text=payload["text"]
80
81 import openai
82 fireworks_client = openai.OpenAI(
83 api_key=os.getenv('FIREWORKS_API_KEY'),
84 base_url="https://api.fireworks.ai/inference/v1"
85 )
86
87 # Check if input text is already in the database
88 existing_documents = [d for d in collection.find({"text": input_text})]
89 document_exists = len(existing_documents) != 0
90
91 # If we already have the document, get the embedding.
92 if document_exists:
93 embedding = existing_documents[0]["embedding"]
94 else:
95 response = fireworks_client.embeddings.create(
96 model="nomic-ai/nomic-embed-text-v1.5",
97 input=input_text,
98 dimensions=768,
99 )
100
101 embedding = response.data[0].embedding
102
103 print(len(embedding))
104
105 # Okay -- if our document exists, check if it has a title, and if it does, use it. Otherwise, generate a title.
106 if document_exists:
107 if "title" in existing_documents[0] and existing_documents[0]["title"] != "":
108 new_title = existing_documents[0]["title"]
109 else:
110 new_title = fire_works_create_title(input_text)
111
112 # Update the document
113 collection.update_one({"_id": existing_documents[0]["_id"]}, {"$set": {"title": new_title}})
114 else:
115 new_title = fire_works_create_title(input_text)
116
117 document = {
118 # "_id": ObjectId("6624177242be2e1d80040517"),
119 "embedding": embedding,
120 "text": input_text,
121 "title": new_title
122 }
123
124 if not document_exists:
125 result = collection.insert_one(document)
126 print("Inserted document ID:", result.inserted_id)
127
128 def query_results(query_vec):
129 results = collection.aggregate([
130 {
131 '$vectorSearch': {
132 "index": "embedding_index",
133 "path": "embedding",
134 "queryVector": query_vec,
135 "numCandidates": 100,
136 "limit": 100,
137 }
138 }
139 ])
140 return results
141
142 # run pipeline
143 results = query_results(embedding)
144 results_collected = [r for r in results]
145 print("reslts collected shape")
146 print(len(results_collected))
147
148 import numpy as np
149 from sklearn.decomposition import PCA
150 from sklearn.preprocessing import MinMaxScaler
151
152 def smush(matrix):
153 pca = PCA(n_components=2)
154 print(matrix.shape)
155 smushed = pca.fit_transform(matrix)
156 return smushed
157
158 # Converts a matrix to a JSON
159 def bundle(texts, titles, matrix):
160 smushed = smush(matrix)
161
162 # Normalize the matrix
163 scaler = MinMaxScaler(feature_range=(-1, 1))
164 smushed = scaler.fit_transform(matrix)
165 smushed = smushed - smushed[0,:]
166
167 return {
168 "data": [
169 {
170 "x": x,
171 "y": y,
172 "text": text,
173 "title": title
174 }
175 for x, y, text, title in zip(smushed[:, 0], smushed[:, 1], texts, titles)
176 ]
177 }
178
179 # Unpack things
180 texts = [t["text"] for t in results_collected]
181 embeddings = [t["embedding"] for t in results_collected]
182
183 print("texts length")
184 print(len(texts))
185
186 print("embeddings length")
187 print(len(embeddings))
188
189 matrix = np.array(embeddings) - embedding
190 return json.dumps(bundle(texts, [""] * len(texts), matrix))
191
192
193def fire_function_call(client):
194
195 messages = [
196 {
197 "role": "system",
198 "content": "You are a very creative assistant with access to functions. Use them if required."
199 },
200 {
201 "role": "user",
202 "content": "Select a super random topic and create a very short title and super detailed long text description for the topic"
203 }
204 ]
205
206 tools = [
207 {
208 "type": "function",
209 "function": {
210 "name": "get_title_text",
211 "description": "Get the random text and title.",
212 "parameters": {
213 "type": "object",
214 "properties": {
215 "random_text": {
216 "type": "string",
217 "description": "random text",
218 },
219 "title": {
220 "type": "string",
221 "description": "title for random text."
222 },
223 },
224 "required": ["random_text", "title"],
225 },
226 },
227 }
228 ]
229
230 chat_completion = client.chat.completions.create(
231 model="accounts/fireworks/models/fw-function-call-34b-v0",
232 messages=messages,
233 tools=tools,
234 tool_choice="auto",
235 temperature=0.8
236 )
237
238 print(repr(chat_completion.choices[0].message.model_dump()))
239
240 message_data = chat_completion.choices[0].message.model_dump()
241
242 # Assuming message_data now contains a dictionary with the 'tool_calls' key
243 tool_calls = message_data['tool_calls']
244
245 # Extract the function arguments from the first tool call
246 function_args_str = tool_calls[0]['function']['arguments']
247
248 # Parse the JSON string in function arguments
249 function_args = json.loads(function_args_str)
250
251 # Extract 'random_text' and 'title' values
252 random_text = function_args['random_text']
253 title = function_args['title']
254
255 print("Random Text:", random_text)
256 print("Title:", title)
257
258 return random_text, title
259
260def fire_works_create_title(text):
261 client = openai.OpenAI(
262 base_url = "https://api.fireworks.ai/inference/v1",
263 api_key=os.getenv('FIREWORKS_API_KEY'),
264 )
265 response = client.chat.completions.create(
266 model="accounts/fireworks/models/llama-v2-7b-chat",
267 messages=[{
268 "role": "user",
269 "content": f"create a title for this text: {text}",
270 }],
271 )
272 print(response.choices[0].message.content)
273
274 return response.choices[0].message.content
275
276
277def fire_function__call_generate_title(client, text):
278
279 messages = [
280 {
281 "role": "system",
282 "content": "You are an assistant good at creative titles with access to functions. Use them if required."
283 },
284 {
285 "role": "user",
286 "content": f"Create a very short title for the input text: {text}"
287 }
288 ]
289
290 tools = [
291 {
292 "type": "function",
293 "function": {
294 "name": "get_title",
295 "description": "Generate a short title.",
296 "parameters": {
297 "type": "object",
298 "properties": {
299 "title": {
300 "type": "string",
301 "description": "title for input text."
302 },
303 },
304 "required": ["title"],
305 },
306 },
307 }
308 ]
309
310 chat_completion = client.chat.completions.create(
311 model="accounts/fireworks/models/fw-function-call-34b-v0",
312 messages=messages,
313 tools=tools,
314 tool_choice="auto",
315 temperature=0.1
316 )
317
318 print(repr(chat_completion.choices[0].message.model_dump()))
319
320 message_data = chat_completion.choices[0].message.model_dump()
321
322 # Assuming message_data now contains a dictionary with the 'tool_calls' key
323 tool_calls = message_data['tool_calls']
324
325 # Extract the function arguments from the first tool call
326 function_args_str = tool_calls[0]['function']['arguments']
327
328 # Parse the JSON string in function arguments
329 function_args = json.loads(function_args_str)
330
331 # Extract 'random_text' and 'title' values
332 title = function_args['title']
333
334 print("Title:", title)
335
336 return title