Upload community_contributions/ederhion_week8/deal_agent_framework.py with huggingface_hub

community_contributions/ederhion_week8/deal_agent_framework.py

@@ -0,0 +1,103 @@
+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
+
+load_dotenv(override=True)
+
+BG_BLUE = "\033[44m"
+WHITE = "\033[37m"
+RESET = "\033[0m"
+
+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"
+    COLLECTION_NAME = "used_cars"
+
+    def __init__(self):
+        init_logging()
+        client = chromadb.PersistentClient(path=self.DB)
+        self.memory = self.read_memory()
+        self.collection = client.get_or_create_collection(self.COLLECTION_NAME)
+        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(cls.COLLECTION_NAME)
+        result = collection.get(
+            include=["embeddings", "documents", "metadatas"], limit=max_datapoints
+        )
+        import numpy as np
+        from sklearn.manifold import TSNE
+        vectors = np.array(result["embeddings"])
+        documents = result["documents"]
+        categories = [metadata.get("category", "Automotive") for metadata in result["metadatas"]]
+        colors = ["blue" for _ 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()