Spaces:
Running on Zero
Running on Zero
File size: 3,897 Bytes
58bd26a | 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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | import os
import sys
import logging
import json
from typing import List
from dotenv import load_dotenv
import chromadb
from agents.planning_agent import PlanningAgent
from agents.deals import Opportunity
from sklearn.manifold import TSNE
import numpy as np
load_dotenv(override=True)
# Colors for logging
BG_BLUE = "\033[44m"
WHITE = "\033[37m"
RESET = "\033[0m"
# Colors for plot
CATEGORIES = [
"Appliances",
"Automotive",
"Cell_Phones_and_Accessories",
"Electronics",
"Musical_Instruments",
"Office_Products",
"Tools_and_Home_Improvement",
"Toys_and_Games",
]
COLORS = ["red", "blue", "brown", "orange", "yellow", "green", "purple", "cyan"]
def init_logging():
root = logging.getLogger()
root.setLevel(logging.INFO)
handler = logging.StreamHandler(sys.stdout)
handler.setLevel(logging.INFO)
formatter = logging.Formatter(
"[%(asctime)s] [Agents] [%(levelname)s] %(message)s",
datefmt="%Y-%m-%d %H:%M:%S %z",
)
handler.setFormatter(formatter)
root.addHandler(handler)
class DealAgentFramework:
DB = "products_vectorstore"
MEMORY_FILENAME = "memory.json"
def __init__(self):
init_logging()
client = chromadb.PersistentClient(path=self.DB)
self.memory = self.read_memory()
self.collection = client.get_or_create_collection("products")
self.planner = None
def init_agents_as_needed(self):
if not self.planner:
self.log("Initializing Agent Framework")
self.planner = PlanningAgent(self.collection)
self.log("Agent Framework is ready")
def read_memory(self) -> List[Opportunity]:
if os.path.exists(self.MEMORY_FILENAME):
with open(self.MEMORY_FILENAME, "r") as file:
data = json.load(file)
opportunities = [Opportunity(**item) for item in data]
return opportunities
return []
def write_memory(self) -> None:
data = [opportunity.model_dump() for opportunity in self.memory]
with open(self.MEMORY_FILENAME, "w") as file:
json.dump(data, file, indent=2)
@classmethod
def reset_memory(cls) -> None:
data = []
if os.path.exists(cls.MEMORY_FILENAME):
with open(cls.MEMORY_FILENAME, "r") as file:
data = json.load(file)
truncated = data[:2]
with open(cls.MEMORY_FILENAME, "w") as file:
json.dump(truncated, file, indent=2)
def log(self, message: str):
text = BG_BLUE + WHITE + "[Agent Framework] " + message + RESET
logging.info(text)
def run(self) -> List[Opportunity]:
self.init_agents_as_needed()
logging.info("Kicking off Planning Agent")
result = self.planner.plan(memory=self.memory)
logging.info(f"Planning Agent has completed and returned: {result}")
if result:
self.memory.append(result)
self.write_memory()
return self.memory
@classmethod
def get_plot_data(cls, max_datapoints=2000):
client = chromadb.PersistentClient(path=cls.DB)
collection = client.get_or_create_collection("products")
result = collection.get(
include=["embeddings", "documents", "metadatas"], limit=max_datapoints
)
vectors = np.array(result["embeddings"])
documents = result["documents"]
categories = [metadata["category"] for metadata in result["metadatas"]]
colors = [COLORS[CATEGORIES.index(c)] for c in categories]
tsne = TSNE(n_components=3, random_state=42, n_jobs=-1)
reduced_vectors = tsne.fit_transform(vectors)
return documents, reduced_vectors, colors
if __name__ == "__main__":
DealAgentFramework().run()
|