refactor
This commit is contained in:
Ryan Chen
251
services/raggr/query.py
Normal file
251
services/raggr/query.py
Normal file
@@ -0,0 +1,251 @@
|
||||
import json
|
||||
import os
|
||||
from typing import Literal
|
||||
import datetime
|
||||
from ollama import Client
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
# Configure ollama client with URL from environment or default to localhost
|
||||
ollama_client = Client(
|
||||
host=os.getenv("OLLAMA_URL", "http://localhost:11434"), timeout=10.0
|
||||
)
|
||||
|
||||
# This uses inferred filters — which means using LLM to create the metadata filters
|
||||
|
||||
|
||||
class FilterOperation(BaseModel):
|
||||
op: Literal["$gt", "$gte", "$eq", "$ne", "$lt", "$lte", "$in", "$nin"]
|
||||
value: str | list[str]
|
||||
|
||||
|
||||
class FilterQuery(BaseModel):
|
||||
field_name: Literal["created_date, tags"]
|
||||
op: FilterOperation
|
||||
|
||||
|
||||
class AndQuery(BaseModel):
|
||||
op: Literal["$and", "$or"]
|
||||
subqueries: list[FilterQuery]
|
||||
|
||||
|
||||
class GeneratedQuery(BaseModel):
|
||||
fields: list[str]
|
||||
extracted_metadata_fields: str
|
||||
|
||||
|
||||
class Time(BaseModel):
|
||||
time: int
|
||||
|
||||
|
||||
DOCTYPE_OPTIONS = [
|
||||
"Bill",
|
||||
"Image Description",
|
||||
"Insurance",
|
||||
"Medical Record",
|
||||
"Documentation",
|
||||
"Letter",
|
||||
]
|
||||
|
||||
QUERY_TYPE_OPTIONS = [
|
||||
"Simba",
|
||||
"Other",
|
||||
]
|
||||
|
||||
|
||||
class DocumentType(BaseModel):
|
||||
type: list[str] = Field(description="type of document", enum=DOCTYPE_OPTIONS)
|
||||
|
||||
|
||||
class QueryType(BaseModel):
|
||||
type: str = Field(desciption="type of query", enum=QUERY_TYPE_OPTIONS)
|
||||
|
||||
|
||||
PROMPT = """
|
||||
You are an information specialist that processes user queries. The current year is 2025. The user queries are all about
|
||||
a cat, Simba, and its records. The types of records are listed below. Using the query, extract the
|
||||
the date range the user is trying to query. You should return it as a JSON. The date tag is created_date. Return the date in epoch time.
|
||||
|
||||
If the created_date cannot be ascertained, set it to epoch time start.
|
||||
|
||||
You have several operators at your disposal:
|
||||
- $gt: greater than
|
||||
- $gte: greater than or equal
|
||||
- $eq: equal
|
||||
- $ne: not equal
|
||||
- $lt: less than
|
||||
- $lte: less than or equal to
|
||||
- $in: in
|
||||
- $nin: not in
|
||||
|
||||
Logical operators:
|
||||
- $and, $or
|
||||
|
||||
### Example 1
|
||||
Query: "Who is Simba's current vet?"
|
||||
Metadata fields: "{"created_date"}"
|
||||
Extracted metadata fields: {"created_date: {"$gt": "2025-01-01"}}
|
||||
|
||||
### Example 2
|
||||
Query: "How many teeth has Simba had removed?"
|
||||
Metadata fields: {}
|
||||
Extracted metadata fields: {}
|
||||
|
||||
### Example 3
|
||||
Query: "How many times has Simba been to the vet this year?"
|
||||
Metadata fields: {"created_date"}
|
||||
Extracted metadata fields: {"created_date": {"gt": "2025-01-01"}}
|
||||
|
||||
document_types:
|
||||
- aftercare
|
||||
- bill
|
||||
- insurance claim
|
||||
- medical records
|
||||
|
||||
Only return the extracted metadata fields. Make sure the extracted metadata fields are valid JSON
|
||||
"""
|
||||
|
||||
|
||||
DOCTYPE_PROMPT = f"""You are an information specialist that processes user queries. A query can have two tags attached from the following options. Based on the query, determine which of the following options is most appropriate: {",".join(DOCTYPE_OPTIONS)}
|
||||
|
||||
### Example 1
|
||||
Query: "Who is Simba's current vet?"
|
||||
Tags: ["Bill", "Medical Record"]
|
||||
|
||||
|
||||
### Example 2
|
||||
Query: "Who does Simba know?"
|
||||
Tags: ["Letter", "Documentation"]
|
||||
"""
|
||||
|
||||
QUERY_TYPE_PROMPT = f"""You are an information specialist that processes user queries.
|
||||
A query can have one tag attached from the following options. Based on the query and the transcript which is listed below, determine
|
||||
which of the following options is most appropriate: {",".join(QUERY_TYPE_OPTIONS)}
|
||||
|
||||
### Example 1
|
||||
Query: "Who is Simba's current vet?"
|
||||
Tags: ["Simba"]
|
||||
|
||||
|
||||
### Example 2
|
||||
Query: "What is the capital of Tokyo?"
|
||||
Tags: ["Other"]
|
||||
|
||||
|
||||
### Example 3
|
||||
Query: "Can you help me write an email?"
|
||||
Tags: ["Other"]
|
||||
|
||||
TRANSCRIPT:
|
||||
"""
|
||||
|
||||
|
||||
class QueryGenerator:
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def date_to_epoch(self, date_str: str) -> float:
|
||||
split_date = date_str.split("-")
|
||||
date = datetime.datetime(
|
||||
int(split_date[0]),
|
||||
int(split_date[1]),
|
||||
int(split_date[2]),
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
)
|
||||
|
||||
return date.timestamp()
|
||||
|
||||
def get_doctype_query(self, input: str):
|
||||
client = OpenAI()
|
||||
response = client.chat.completions.create(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are an information specialist that is really good at deciding what tags a query should have",
|
||||
},
|
||||
{"role": "user", "content": DOCTYPE_PROMPT + " " + input},
|
||||
],
|
||||
model="gpt-4o",
|
||||
response_format={
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "document_type",
|
||||
"schema": DocumentType.model_json_schema(),
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
response_json_str = response.choices[0].message.content
|
||||
type_data = json.loads(response_json_str)
|
||||
metadata_query = {"document_type": {"$in": type_data["type"]}}
|
||||
return metadata_query
|
||||
|
||||
def get_query_type(self, input: str, transcript: str):
|
||||
client = OpenAI()
|
||||
response = client.chat.completions.create(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are an information specialist that is really good at deciding what tags a query should have",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"{QUERY_TYPE_PROMPT}\nTRANSCRIPT:\n{transcript}\nQUERY:{input}",
|
||||
},
|
||||
],
|
||||
model="gpt-4o",
|
||||
response_format={
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "query_type",
|
||||
"schema": QueryType.model_json_schema(),
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
response_json_str = response.choices[0].message.content
|
||||
type_data = json.loads(response_json_str)
|
||||
return type_data["type"]
|
||||
|
||||
def get_query(self, input: str):
|
||||
client = OpenAI()
|
||||
response = client.responses.parse(
|
||||
model="gpt-4o",
|
||||
input=[
|
||||
{"role": "system", "content": PROMPT},
|
||||
{"role": "user", "content": input},
|
||||
],
|
||||
text_format=GeneratedQuery,
|
||||
)
|
||||
print(response.output)
|
||||
query = json.loads(response.output_parsed.extracted_metadata_fields)
|
||||
# response: ChatResponse = ollama_client.chat(
|
||||
# model="gemma3n:e4b",
|
||||
# messages=[
|
||||
# {"role": "system", "content": PROMPT},
|
||||
# {"role": "user", "content": input},
|
||||
# ],
|
||||
# format=GeneratedQuery.model_json_schema(),
|
||||
# )
|
||||
|
||||
# query = json.loads(
|
||||
# json.loads(response["message"]["content"])["extracted_metadata_fields"]
|
||||
# )
|
||||
# date_key = list(query["created_date"].keys())[0]
|
||||
# query["created_date"][date_key] = self.date_to_epoch(
|
||||
# query["created_date"][date_key]
|
||||
# )
|
||||
|
||||
# if "$" not in date_key:
|
||||
# query["created_date"]["$" + date_key] = query["created_date"][date_key]
|
||||
|
||||
return query
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
qg = QueryGenerator()
|
||||
print(qg.get_doctype_query("How heavy is Simba?"))
|
||||
Reference in New Issue
Block a user