105 lines
3.9 KiB
Python
105 lines
3.9 KiB
Python
# imports
|
|
|
|
import os
|
|
import re
|
|
import math
|
|
import json
|
|
from typing import List, Dict
|
|
from sentence_transformers import SentenceTransformer
|
|
from datasets import load_dataset
|
|
import chromadb
|
|
from items import Item
|
|
from testing import Tester
|
|
from agents.agent import Agent
|
|
from groq import Groq # Only Groq is used
|
|
|
|
|
|
class FrontierAgent(Agent):
|
|
|
|
name = "Frontier Agent"
|
|
color = Agent.BLUE
|
|
|
|
MODEL = "meta-llama/llama-4-scout-17b-16e-instruct"
|
|
|
|
def __init__(self, collection):
|
|
"""
|
|
Set up this instance by connecting to Groq,
|
|
connect to the Chroma Datastore, and set up the vector encoding model.
|
|
"""
|
|
self.log("Initializing Frontier Agent")
|
|
|
|
groq_api_key = os.getenv("GROQ_API_KEY")
|
|
if not groq_api_key:
|
|
raise ValueError("GROQ_API_KEY environment variable not set.")
|
|
|
|
self.client = Groq(api_key=groq_api_key)
|
|
self.log("Frontier Agent is set up with Groq")
|
|
|
|
self.collection = collection
|
|
self.model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
|
|
self.log("Frontier Agent is ready")
|
|
|
|
def make_context(self, similars: List[str], prices: List[float]) -> str:
|
|
"""
|
|
Create context that can be inserted into the prompt
|
|
"""
|
|
message = "To provide some context, here are some other items that might be similar to the item you need to estimate.\n\n"
|
|
for similar, price in zip(similars, prices):
|
|
message += f"Potentially related product:\n{similar}\nPrice is ${price:.2f}\n\n"
|
|
return message
|
|
|
|
def messages_for(self, description: str, similars: List[str], prices: List[float]) -> List[Dict[str, str]]:
|
|
"""
|
|
Create the message list to be included in a call to the language model
|
|
"""
|
|
system_message = "You estimate prices of items. Reply only with the price, no explanation"
|
|
user_prompt = self.make_context(similars, prices)
|
|
user_prompt += "And now the question for you:\n\n"
|
|
user_prompt += "How much does this cost?\n\n" + description
|
|
return [
|
|
{"role": "system", "content": system_message},
|
|
{"role": "user", "content": user_prompt},
|
|
{"role": "assistant", "content": "Price is $"}
|
|
]
|
|
|
|
def find_similars(self, description: str):
|
|
"""
|
|
Return a list of items similar to the given one by looking in the Chroma datastore
|
|
"""
|
|
self.log("Frontier Agent is performing a RAG search of the Chroma datastore to find 5 similar products")
|
|
vector = self.model.encode([description])
|
|
results = self.collection.query(query_embeddings=vector.astype(float).tolist(), n_results=5)
|
|
documents = results['documents'][0][:]
|
|
prices = [m['price'] for m in results['metadatas'][0][:]]
|
|
self.log("Frontier Agent has found similar products")
|
|
return documents, prices
|
|
|
|
def get_price(self, s) -> float:
|
|
"""
|
|
A utility that plucks a floating point number out of a string
|
|
"""
|
|
s = s.replace('$', '').replace(',', '')
|
|
match = re.search(r"[-+]?\d*\.\d+|\d+", s)
|
|
return float(match.group()) if match else 0.0
|
|
|
|
def price(self, description: str) -> float:
|
|
"""
|
|
Make a call to Groq to estimate the price of the described product,
|
|
by looking up 5 similar products and including them in the prompt to give context
|
|
"""
|
|
documents, prices = self.find_similars(description)
|
|
self.log(f"Frontier Agent is calling {self.MODEL} via Groq with context including 5 similar products")
|
|
|
|
response = self.client.chat.completions.create(
|
|
model=self.MODEL,
|
|
messages=self.messages_for(description, documents, prices),
|
|
temperature=0.0,
|
|
max_tokens=5,
|
|
seed=42
|
|
)
|
|
|
|
reply = response.choices[0].message.content
|
|
result = self.get_price(reply)
|
|
self.log(f"Frontier Agent completed - predicting ${result:.2f}")
|
|
return result
|