-
Notifications
You must be signed in to change notification settings - Fork 0
/
query.py
71 lines (67 loc) · 1.98 KB
/
query.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
import polars as pl
import os
from qdrant_client import models, QdrantClient
from sentence_transformers import SentenceTransformer
from openai import OpenAI
import textwrap
from dotenv import load_dotenv
load_dotenv()
os.environ["TOKENIZERS_PARALLELISM"] = "true"
encoder = SentenceTransformer("all-MiniLM-L6-v2")
qdrant = QdrantClient("http://localhost:6333")
client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
PROMPT = input("Enter a query: ")
encoded = encoder.encode(PROMPT).tolist()
hits = qdrant.search(
collection_name="recipes",
query_vector=models.NamedVector(name="title", vector=encoded),
query_filter=models.Filter(
must=[
models.FieldCondition(
key="likes",
range=models.Range(gte=10),
)
]
),
limit=5,
)
df = pl.DataFrame(
[
{
"title": hit.payload["title"],
"ingredients": hit.payload["ingredients"],
"steps": hit.payload["steps"],
"score": hit.score,
}
for hit in hits
]
)
print(df)
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=[
{
"role": "system",
"content": textwrap.dedent("""
You are a chef chatbot.
Your top priority is to help users into picking one recipe that suits the user request the best.
If no recipe suits the user request perfectly, pick the one that is the closest.
Recipes are in Indonesian, automatically translate the recipe to English.
Just give the name of the dish, the ingredients, and the steps to make the dish in English.
"""),
},
{
"role": "user",
"content": PROMPT,
},
{
"role": "assistant",
"content": str(df.to_dicts()),
},
],
stream=True,
)
for chunk in response:
if chunk.choices[0].delta.content:
print(chunk.choices[0].delta.content, end="")
print()