Spaces:
Running on Zero
Running on Zero
| import os | |
| import re | |
| from typing import List, Dict | |
| from groq import Groq | |
| from sentence_transformers import SentenceTransformer | |
| from agents.agent import Agent | |
| from free_config import GROQ_MODEL | |
| class FrontierAgent(Agent): | |
| name = "Frontier Agent" | |
| color = Agent.BLUE | |
| def __init__(self, collection): | |
| """ | |
| Set up this instance by connecting to Groq, the Chroma datastore, | |
| and the local sentence-transformer embedding model. | |
| """ | |
| self.log("Initializing Frontier Agent") | |
| self.client = Groq(api_key=os.environ["GROQ_API_KEY"]) | |
| self.model = GROQ_MODEL | |
| self.log(f"Frontier Agent is setting up with Groq ({self.model})") | |
| self.collection = collection | |
| self.encoder = 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 for the Groq chat completion call | |
| """ | |
| message = f"Estimate the price of this product. Respond with the price only, no explanation.\n\n{description}\n\n" | |
| message += self.make_context(similars, prices) | |
| return [{"role": "user", "content": message}] | |
| 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.encoder.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: | |
| """ | |
| Call Groq to estimate the price using RAG context from similar products. | |
| """ | |
| documents, prices = self.find_similars(description) | |
| self.log( | |
| f"Frontier Agent is about to call Groq ({self.model}) with context including 5 similar products" | |
| ) | |
| response = self.client.chat.completions.create( | |
| model=self.model, | |
| messages=self.messages_for(description, documents, prices), | |
| temperature=0, | |
| ) | |
| reply = response.choices[0].message.content | |
| result = self.get_price(reply) | |
| self.log(f"Frontier Agent completed - predicting ${result:.2f}") | |
| return result | |