From e044f9c639f7c29072f5dc1f3fbd862716bdd7ad Mon Sep 17 00:00:00 2001
From: James Ketrenos <james_git@ketrenos.com>
Date: Wed, 14 May 2025 11:31:31 -0700
Subject: [PATCH] Reformatted all content to black

---
 frontend/src/App.css                |    4 +
 frontend/src/BackstoryTextField.tsx |    1 -
 frontend/src/ChatBubble.tsx         |    2 +-
 frontend/src/Conversation.tsx       |  115 +-
 frontend/src/Message.tsx            |   74 +-
 frontend/src/ResumeBuilderPage.tsx  |   31 +-
 frontend/src/StyledMarkdown.css     |    9 +
 frontend/src/StyledMarkdown.tsx     |   10 +-
 src/server.py                       |  587 ++++---
 src/tests/test-agent.py             |   22 +-
 src/tests/test-context-routing.py   |   86 +-
 src/tests/test-context.py           |   18 +-
 src/tests/test-message.py           |    8 +-
 src/tests/test-metrics.py           |    8 +-
 src/utils/__init__.py               |   59 +-
 src/utils/agents/__init__.py        |   14 +-
 src/utils/agents/base.py            |  302 ++--
 src/utils/agents/chat.py            |   56 +-
 src/utils/agents/fact_check.py      |   88 +-
 src/utils/agents/job_description.py | 2313 +++++++++++++++++----------
 src/utils/agents/resume.py          |  150 +-
 src/utils/agents/types.py           |   51 +-
 src/utils/chroma.py                 |  122 --
 src/utils/chunk.py                  |   88 -
 src/utils/context.py                |  115 +-
 src/utils/conversation.py           |   23 +-
 src/utils/defines.py                |   16 +-
 src/utils/full.py                   |  468 ------
 src/utils/message.py                |   74 +-
 src/utils/metrics.py                |  156 +-
 src/utils/rag.py                    |  504 +++---
 src/utils/setup_logging.py          |   23 +-
 src/utils/tools/__init__.py         |    8 +-
 src/utils/tools/basetools.py        |  371 +++--
 34 files changed, 3223 insertions(+), 2753 deletions(-)
 delete mode 100644 src/utils/chroma.py
 delete mode 100644 src/utils/chunk.py
 delete mode 100644 src/utils/full.py

diff --git a/frontend/src/App.css b/frontend/src/App.css
index 4331fa7..4928e12 100644
--- a/frontend/src/App.css
+++ b/frontend/src/App.css
@@ -88,6 +88,10 @@ button {
   flex-grow: 1;
 }
 
+.MessageContent div > p:first-child {
+  margin-top: 0;
+}
+
 .MenuCard.MuiCard-root {
   display: flex;
   flex-direction: column;
diff --git a/frontend/src/BackstoryTextField.tsx b/frontend/src/BackstoryTextField.tsx
index edddfb8..9046a42 100644
--- a/frontend/src/BackstoryTextField.tsx
+++ b/frontend/src/BackstoryTextField.tsx
@@ -30,7 +30,6 @@ const BackstoryTextField = React.forwardRef<BackstoryTextFieldRef, BackstoryText
   const shadowRef = useRef<HTMLTextAreaElement>(null);
   const [editValue, setEditValue] = useState<string>(value);
 
-  console.log({ value, placeholder, editValue });
   // Sync editValue with prop value if it changes externally
   useEffect(() => {
     setEditValue(value || "");
diff --git a/frontend/src/ChatBubble.tsx b/frontend/src/ChatBubble.tsx
index bb82895..80be313 100644
--- a/frontend/src/ChatBubble.tsx
+++ b/frontend/src/ChatBubble.tsx
@@ -158,7 +158,7 @@ function ChatBubble(props: ChatBubbleProps) {
   };
 
   // Render Accordion for expandable content
-  if (expandable || (role === 'content' && title)) {
+  if (expandable || title) {
     // Determine if Accordion is controlled
     const isControlled = typeof expanded === 'boolean' && typeof onExpand === 'function';
 
diff --git a/frontend/src/Conversation.tsx b/frontend/src/Conversation.tsx
index f21dfdb..91e4a92 100644
--- a/frontend/src/Conversation.tsx
+++ b/frontend/src/Conversation.tsx
@@ -1,13 +1,15 @@
 import React, { useState, useImperativeHandle, forwardRef, useEffect, useRef, useCallback } from 'react';
 import Typography from '@mui/material/Typography';
 import Tooltip from '@mui/material/Tooltip';
+import IconButton from '@mui/material/IconButton';
 import Button from '@mui/material/Button';
 import Box from '@mui/material/Box';
 import SendIcon from '@mui/icons-material/Send';
+import CancelIcon from '@mui/icons-material/Cancel';
 import { SxProps, Theme } from '@mui/material';
 import PropagateLoader from "react-spinners/PropagateLoader";
 
-import { Message, MessageList, MessageData } from './Message';
+import { Message, MessageList, BackstoryMessage } from './Message';
 import { ContextStatus } from './ContextStatus';
 import { Scrollable } from './Scrollable';
 import { DeleteConfirmation } from './DeleteConfirmation';
@@ -17,7 +19,7 @@ import { BackstoryTextField, BackstoryTextFieldRef } from './BackstoryTextField'
 import { BackstoryElementProps } from './BackstoryTab';
 import { connectionBase } from './Global';
 
-const loadingMessage: MessageData = { "role": "status", "content": "Establishing connection with server..." };
+const loadingMessage: BackstoryMessage = { "role": "status", "content": "Establishing connection with server..." };
 
 type ConversationMode = 'chat' | 'job_description' | 'resume' | 'fact_check';
 
@@ -25,24 +27,6 @@ interface ConversationHandle {
   submitQuery: (prompt: string, options?: QueryOptions) => void;
   fetchHistory: () => void;
 }
-interface BackstoryMessage {
-  prompt: string;
-  preamble: {};
-  status: string;
-  full_content: string;
-  response: string;     // Set when status === 'done' or 'error'
-  chunk: string;        // Used when status === 'streaming'
-  metadata: {
-    rag: { documents: [] };
-    tools: string[];
-    eval_count: number;
-    eval_duration: number;
-    prompt_eval_count: number;
-    prompt_eval_duration: number;
-  };
-  actions: string[];
-  timestamp: string;
-};
 
 interface ConversationProps extends BackstoryElementProps {
   className?: string,         // Override default className
@@ -59,7 +43,7 @@ interface ConversationProps extends BackstoryElementProps {
   messageFilter?: ((messages: MessageList) => MessageList) | undefined, // Filter callback to determine which Messages to display in Conversation
   messages?: MessageList,     // 
   sx?: SxProps<Theme>,
-  onResponse?: ((message: MessageData) => void) | undefined, // Event called when a query completes (provides messages)
+  onResponse?: ((message: BackstoryMessage) => void) | undefined, // Event called when a query completes (provides messages)
 };
 
 const Conversation = forwardRef<ConversationHandle, ConversationProps>(({
@@ -87,8 +71,8 @@ const Conversation = forwardRef<ConversationHandle, ConversationProps>(({
   const [countdown, setCountdown] = useState<number>(0);
   const [conversation, setConversation] = useState<MessageList>([]);
   const [filteredConversation, setFilteredConversation] = useState<MessageList>([]);
-  const [processingMessage, setProcessingMessage] = useState<MessageData | undefined>(undefined);
-  const [streamingMessage, setStreamingMessage] = useState<MessageData | undefined>(undefined);
+  const [processingMessage, setProcessingMessage] = useState<BackstoryMessage | undefined>(undefined);
+  const [streamingMessage, setStreamingMessage] = useState<BackstoryMessage | undefined>(undefined);
   const timerRef = useRef<any>(null);
   const [contextStatus, setContextStatus] = useState<ContextStatus>({ context_used: 0, max_context: 0 });
   const [contextWarningShown, setContextWarningShown] = useState<boolean>(false);
@@ -96,6 +80,7 @@ const Conversation = forwardRef<ConversationHandle, ConversationProps>(({
   const conversationRef = useRef<MessageList>([]);
   const viewableElementRef = useRef<HTMLDivElement>(null);
   const backstoryTextRef = useRef<BackstoryTextFieldRef>(null);
+  const stopRef = useRef(false);
 
   // Keep the ref updated whenever items changes
   useEffect(() => {
@@ -181,14 +166,25 @@ const Conversation = forwardRef<ConversationHandle, ConversationProps>(({
 
           const backstoryMessages: BackstoryMessage[] = messages;
 
-          setConversation(backstoryMessages.flatMap((backstoryMessage: BackstoryMessage) => [{
-            role: 'user',
-            content: backstoryMessage.prompt || "",
-          }, {
-            ...backstoryMessage,
-              role: backstoryMessage.status === "done" ? "assistant" : backstoryMessage.status,
+          setConversation(backstoryMessages.flatMap((backstoryMessage: BackstoryMessage) => {
+            if (backstoryMessage.status === "partial") {
+              return [{
+                ...backstoryMessage,
+                role: "assistant",
+                content: backstoryMessage.response || "",
+                expanded: false,
+                expandable: true,
+              }]
+            }
+            return [{
+              role: 'user',
+              content: backstoryMessage.prompt || "",
+            }, {
+                ...backstoryMessage,
+              role: ['done'].includes(backstoryMessage.status || "") ? "assistant" : backstoryMessage.status,
               content: backstoryMessage.response || "",
-            }] as MessageList));
+              }] as MessageList;
+          }));
           setNoInteractions(false);
         }
         setProcessingMessage(undefined);
@@ -294,6 +290,11 @@ const Conversation = forwardRef<ConversationHandle, ConversationProps>(({
     }
   };
 
+  const cancelQuery = () => {
+    console.log("Stop query");
+    stopRef.current = true;
+  };
+
   const sendQuery = async (request: string, options?: QueryOptions) => {
     request = request.trim();
 
@@ -308,6 +309,8 @@ const Conversation = forwardRef<ConversationHandle, ConversationProps>(({
       return;
     }
 
+    stopRef.current = false;
+
     setNoInteractions(false);
 
     setConversation([
@@ -325,12 +328,10 @@ const Conversation = forwardRef<ConversationHandle, ConversationProps>(({
 
     try {
       setProcessing(true);
-      // Create a unique ID for the processing message
-      const processingId = Date.now().toString();
 
       // Add initial processing message
       setProcessingMessage(
-        { role: 'status', content: 'Submitting request...', id: processingId, isProcessing: true }
+        { role: 'status', content: 'Submitting request...', disableCopy: true }
       );
 
       // Add a small delay to ensure React has time to update the UI
@@ -379,17 +380,20 @@ const Conversation = forwardRef<ConversationHandle, ConversationProps>(({
 
         switch (update.status) {
           case 'done':
-            console.log('Done processing:', update);
-            stopCountdown();
-            setStreamingMessage(undefined);
-            setProcessingMessage(undefined);
+          case 'partial':
+            if (update.status === 'done') stopCountdown();
+            if (update.status === 'done') setStreamingMessage(undefined);
+            if (update.status === 'done') setProcessingMessage(undefined);
             const backstoryMessage: BackstoryMessage = update;
             setConversation([
               ...conversationRef.current, {
                 ...backstoryMessage,
                 role: 'assistant',
                 origin: type,
+                prompt: ['done', 'partial'].includes(update.status) ? update.prompt : '',
                 content: backstoryMessage.response || "",
+                expanded: update.status === "done" ? true : false,
+                expandable: true,
               }] as MessageList);
             // Add a small delay to ensure React has time to update the UI
             await new Promise(resolve => setTimeout(resolve, 0));
@@ -424,9 +428,9 @@ const Conversation = forwardRef<ConversationHandle, ConversationProps>(({
             // Update processing message with immediate re-render
             if (update.status === "streaming") {
               streaming_response += update.chunk
-              setStreamingMessage({ role: update.status, content: streaming_response });
+              setStreamingMessage({ role: update.status, content: streaming_response, disableCopy: true });
             } else {
-              setProcessingMessage({ role: update.status, content: update.response });
+              setProcessingMessage({ role: update.status, content: update.response, disableCopy: true });
               /* Reset stream on non streaming message */
               streaming_response = ""
             }
@@ -437,12 +441,11 @@ const Conversation = forwardRef<ConversationHandle, ConversationProps>(({
         }
       }
 
-      while (true) {
+      while (!stopRef.current) {
         const { done, value } = await reader.read();
         if (done) {
           break;
         }
-
         const chunk = decoder.decode(value, { stream: true });
 
         // Process each complete line immediately
@@ -470,26 +473,32 @@ const Conversation = forwardRef<ConversationHandle, ConversationProps>(({
         }
       }
 
+      if (stopRef.current) {
+        await reader.cancel();
+        setProcessingMessage(undefined);
+        setStreamingMessage(undefined);
+        setSnack("Processing cancelled", "warning");
+      }
       stopCountdown();
       setProcessing(false);
+      stopRef.current = false;
     } catch (error) {
       console.error('Fetch error:', error);
       setSnack("Unable to process query", "error");
-      setProcessingMessage({ role: 'error', content: "Unable to process query" });
+      setProcessingMessage({ role: 'error', content: "Unable to process query", disableCopy: true });
       setTimeout(() => {
         setProcessingMessage(undefined);
       }, 5000);
-
+      stopRef.current = false;
       setProcessing(false);
       stopCountdown();
-      // Add a small delay to ensure React has time to update the UI
-      await new Promise(resolve => setTimeout(resolve, 0));
+      return;
     }
   };
 
   return (
     <Scrollable
-      className={className || "Conversation"}
+      className={`${className || ""} Conversation`}
       autoscroll
       textFieldRef={viewableElementRef}
       fallbackThreshold={0.5}
@@ -564,6 +573,20 @@ const Conversation = forwardRef<ConversationHandle, ConversationProps>(({
               </Button>
             </span>
           </Tooltip>
+          <Tooltip title="Cancel">
+            <span style={{ display: "flex" }}> { /* This span is used to wrap the IconButton to ensure Tooltip works even when disabled */}
+              <IconButton
+                aria-label="cancel"
+                onClick={() => { cancelQuery(); }}
+                sx={{ display: "flex", margin: 'auto 0px' }}
+                size="large"
+                edge="start"
+                disabled={stopRef.current || sessionId === undefined || processing === false}
+              >
+                <CancelIcon />
+              </IconButton>
+            </span>
+          </Tooltip>          
         </Box>
       </Box>
       {(noInteractions || !hideDefaultPrompts) && defaultPrompts !== undefined && defaultPrompts.length &&
diff --git a/frontend/src/Message.tsx b/frontend/src/Message.tsx
index 74b01cf..517d1d7 100644
--- a/frontend/src/Message.tsx
+++ b/frontend/src/Message.tsx
@@ -47,11 +47,18 @@ type MessageRoles =
   'thinking' |
   'user';
 
-type MessageData = {
+type BackstoryMessage = {
+  // Only two required fields
   role: MessageRoles,
   content: string,
-  status?: string, // streaming, done, error...
-  response?: string,
+  // Rest are optional
+  prompt?: string;
+  preamble?: {};
+  status?: string;
+  full_content?: string;
+  response?: string;     // Set when status === 'done', 'partial', or 'error'
+  chunk?: string;        // Used when status === 'streaming'
+  timestamp?: string;
   disableCopy?: boolean,
   user?: string,
   title?: string,
@@ -84,11 +91,11 @@ interface MessageMetaData {
   setSnack: SetSnackType,
 }
 
-type MessageList = MessageData[];
+type MessageList = BackstoryMessage[];
 
 interface MessageProps extends BackstoryElementProps {
   sx?: SxProps<Theme>,
-  message: MessageData,
+  message: BackstoryMessage,
   expanded?: boolean,
   onExpand?: (open: boolean) => void,
   className?: string,
@@ -237,7 +244,7 @@ const MessageMeta = (props: MessageMetaProps) => {
 };
 
 const Message = (props: MessageProps) => {
-  const { message, submitQuery, sx, className, onExpand, expanded, sessionId, setSnack } = props;
+  const { message, submitQuery, sx, className, onExpand, sessionId, setSnack } = props;
   const [metaExpanded, setMetaExpanded] = useState<boolean>(false);
   const textFieldRef = useRef(null);
 
@@ -254,14 +261,16 @@ const Message = (props: MessageProps) => {
     return (<></>);
   }
 
-  const formattedContent = message.content.trim() || "Waiting for LLM to spool up...";
+  const formattedContent = message.content.trim();
+  if (formattedContent === "") {
+    return (<></>);
+  }
 
   return (
     <ChatBubble
-      className={className || "Message"}
+      className={`${className || ""} Message Message-${message.role}`}
       {...message}
       onExpand={onExpand}
-      expanded={expanded}
       sx={{
         display: "flex",
         flexDirection: "column",
@@ -273,34 +282,24 @@ const Message = (props: MessageProps) => {
         ...sx,
       }}>
       <CardContent ref={textFieldRef} sx={{ position: "relative", display: "flex", flexDirection: "column", overflowX: "auto", m: 0, p: 0, paddingBottom: '0px !important' }}>
-        {message.role !== 'user' ?
-          <Scrollable
-            className="MessageContent"
-            autoscroll
-            fallbackThreshold={0.5}
-            sx={{
-              p: 0,
-              m: 0,
-              maxHeight: (message.role === "streaming") ? "20rem" : "unset",
-              display: "flex",
-              flexGrow: 1,
-              overflow: "auto", /* Handles scrolling for the div */
-            }}
-          >
-            <StyledMarkdown streaming={message.role === "streaming"} {...{ content: formattedContent, submitQuery, sessionId, setSnack }} />
-          </Scrollable>
-           :
-          <Typography
-            className="MessageContent"
-            ref={textFieldRef}
-            variant="body2"
-            sx={{ display: "flex", color: 'text.secondary' }}>
-            {message.content}
-          </Typography>
-        }
+        <Scrollable
+          className="MessageContent"
+          autoscroll
+          fallbackThreshold={0.5}
+          sx={{
+            p: 0,
+            m: 0,
+            maxHeight: (message.role === "streaming") ? "20rem" : "unset",
+            display: "flex",
+            flexGrow: 1,
+            overflow: "auto", /* Handles scrolling for the div */
+          }}
+        >
+          <StyledMarkdown streaming={message.role === "streaming"} {...{ content: formattedContent, submitQuery, sessionId, setSnack }} />
+        </Scrollable>
       </CardContent>
       <CardActions disableSpacing sx={{ display: "flex", flexDirection: "row", justifyContent: "space-between", alignItems: "center", width: "100%", p: 0, m: 0 }}>
-        {(message.disableCopy === undefined || message.disableCopy === false) && ["assistant", "content"].includes(message.role) && <CopyBubble content={message.content} />}
+        {(message.disableCopy === undefined || message.disableCopy === false) && <CopyBubble content={message.content} />}
         {message.metadata && (
           <Box sx={{ display: "flex", alignItems: "center", gap: 1 }}>
             <Button variant="text" onClick={handleMetaExpandClick} sx={{ color: "darkgrey", p: 0 }}>
@@ -309,7 +308,7 @@ const Message = (props: MessageProps) => {
             <ExpandMore
               expand={metaExpanded}
               onClick={handleMetaExpandClick}
-              aria-expanded={expanded}
+              aria-expanded={message.expanded}
               aria-label="show more"
             >
               <ExpandMoreIcon />
@@ -331,7 +330,8 @@ const Message = (props: MessageProps) => {
 export type {
   MessageProps, 
   MessageList,
-  MessageData,
+  BackstoryMessage,
+  MessageMetaData,
   MessageRoles,
 };
 
diff --git a/frontend/src/ResumeBuilderPage.tsx b/frontend/src/ResumeBuilderPage.tsx
index 6fdbc91..0d9d1a5 100644
--- a/frontend/src/ResumeBuilderPage.tsx
+++ b/frontend/src/ResumeBuilderPage.tsx
@@ -7,7 +7,7 @@ import {
 import { SxProps } from '@mui/material';
 
 import { ChatQuery } from './ChatQuery';
-import { MessageList, MessageData } from './Message';
+import { MessageList, BackstoryMessage } from './Message';
 import { Conversation } from './Conversation';
 import { BackstoryPageProps } from './BackstoryTab';
 
@@ -62,11 +62,6 @@ const ResumeBuilderPage: React.FC<BackstoryPageProps> = ({
       return [];
     }
 
-    if (messages.length > 2) {
-      setHasResume(true);
-      setHasFacts(true);
-    }
-
     if (messages.length > 0) {
       messages[0].role = 'content';
       messages[0].title = 'Job Description';
@@ -74,6 +69,19 @@ const ResumeBuilderPage: React.FC<BackstoryPageProps> = ({
       messages[0].expandable = true;
     }
 
+    if (-1 !== messages.findIndex(m => m.status === 'done')) {
+      setHasResume(true);
+      setHasFacts(true);
+    }
+
+    return messages;
+
+    if (messages.length > 1) {
+      setHasResume(true);
+      setHasFacts(true);
+    }
+
+
     if (messages.length > 3) {
       // messages[2] is Show job requirements
       messages[3].role = 'job-requirements';
@@ -95,6 +103,8 @@ const ResumeBuilderPage: React.FC<BackstoryPageProps> = ({
       return [];
     }
 
+    return messages;
+
     if (messages.length > 1) {
       // messages[0] is Show Qualifications
       messages[1].role = 'qualifications';
@@ -139,7 +149,7 @@ const ResumeBuilderPage: React.FC<BackstoryPageProps> = ({
     return filtered;
   }, []);
 
-  const jobResponse = useCallback(async (message: MessageData) => {
+  const jobResponse = useCallback(async (message: BackstoryMessage) => {
     console.log('onJobResponse', message);
     if (message.actions && message.actions.includes("job_description")) {
       await jobConversationRef.current.fetchHistory();
@@ -155,12 +165,12 @@ const ResumeBuilderPage: React.FC<BackstoryPageProps> = ({
     }
   }, [setHasFacts, setHasResume, setActiveTab]);
 
-  const resumeResponse = useCallback((message: MessageData): void => {
+  const resumeResponse = useCallback((message: BackstoryMessage): void => {
     console.log('onResumeResponse', message);
     setHasFacts(true);
   }, [setHasFacts]);
 
-  const factsResponse = useCallback((message: MessageData): void => {
+  const factsResponse = useCallback((message: BackstoryMessage): void => {
     console.log('onFactsResponse', message);
   }, []);
 
@@ -199,7 +209,8 @@ const ResumeBuilderPage: React.FC<BackstoryPageProps> = ({
   3. **Mapping Analysis**: Identifies legitimate matches between requirements and qualifications
 3. **Resume Generation**: Uses mapping output to create a tailored resume with evidence-based content
 4. **Verification**: Performs fact-checking to catch any remaining fabrications
-  1. **Re-generation**: If verification does not pass, a second attempt is made to correct any issues`
+  1. **Re-generation**: If verification does not pass, a second attempt is made to correct any issues`,
+      disableCopy: true
     }];
 
 
diff --git a/frontend/src/StyledMarkdown.css b/frontend/src/StyledMarkdown.css
index 576a857..79af80f 100644
--- a/frontend/src/StyledMarkdown.css
+++ b/frontend/src/StyledMarkdown.css
@@ -16,4 +16,13 @@ pre:not(.MessageContent) {
 
 .MuiMarkdown > div {
   width: 100%;
+}
+
+.Message-streaming .MuiMarkdown ul,
+.Message-streaming .MuiMarkdown h1,
+.Message-streaming .MuiMarkdown p,
+.Message-assistant .MuiMarkdown ul,
+.Message-assistant .MuiMarkdown h1,
+.Message-assistant .MuiMarkdown p {
+  color: white;
 }
\ No newline at end of file
diff --git a/frontend/src/StyledMarkdown.tsx b/frontend/src/StyledMarkdown.tsx
index afb636d..e862044 100644
--- a/frontend/src/StyledMarkdown.tsx
+++ b/frontend/src/StyledMarkdown.tsx
@@ -37,14 +37,14 @@ const StyledMarkdown: React.FC<StyledMarkdownProps> = (props: StyledMarkdownProp
         }
         if (className === "lang-json" && !streaming) {
           try {
-            const fixed = jsonrepair(content);
+            let fixed = JSON.parse(jsonrepair(content));
             return <Scrollable className="JsonViewScrollable">
               <JsonView
                 className="JsonView"
                 style={{
                   ...vscodeTheme,
                   fontSize: "0.8rem",
-                  maxHeight: "20rem",
+                  maxHeight: "10rem",
                   padding: "14px 0",
                   overflow: "hidden",
                   width: "100%",
@@ -53,9 +53,9 @@ const StyledMarkdown: React.FC<StyledMarkdownProps> = (props: StyledMarkdownProp
                 }}
                 displayDataTypes={false}
                 objectSortKeys={false}
-                collapsed={false}
+                collapsed={true}
                 shortenTextAfterLength={100}
-                value={JSON.parse(fixed)}>
+                value={fixed}>
                 <JsonView.String
                   render={({ children, ...reset }) => {
                     if (typeof (children) === "string" && children.match("\n")) {
@@ -66,7 +66,7 @@ const StyledMarkdown: React.FC<StyledMarkdownProps> = (props: StyledMarkdownProp
               </JsonView>
             </Scrollable>
           } catch (e) {
-            console.log("jsonrepair error", e);
+            return <pre><code className="JsonRaw">{content}</code></pre>
           };
         }
         return <pre><code className={className}>{element.children}</code></pre>;
diff --git a/src/server.py b/src/server.py
index ab73693..a97cea7 100644
--- a/src/server.py
+++ b/src/server.py
@@ -1,8 +1,9 @@
-LLM_TIMEOUT=600
+LLM_TIMEOUT = 600
 
 from utils import logger
+from pydantic import BaseModel, Field  # type: ignore
 
-from typing import AsyncGenerator
+from typing import AsyncGenerator, Dict, Optional
 
 # %%
 # Imports [standard]
@@ -26,6 +27,7 @@ from uuid import uuid4
 import time
 import traceback
 
+
 def try_import(module_name, pip_name=None):
     try:
         __import__(module_name)
@@ -33,6 +35,7 @@ def try_import(module_name, pip_name=None):
         print(f"Module '{module_name}' not found. Install it using:")
         print(f"  pip install {pip_name or module_name}")
 
+
 # Third-party modules with import checks
 try_import("ollama")
 try_import("requests")
@@ -47,23 +50,25 @@ try_import("prometheus_fastapi_instrumentator")
 import ollama
 import requests
 from contextlib import asynccontextmanager
-from fastapi import FastAPI, Request, BackgroundTasks # type: ignore
-from fastapi.responses import JSONResponse, StreamingResponse, FileResponse, RedirectResponse # type: ignore
-from fastapi.middleware.cors import CORSMiddleware # type: ignore
-import uvicorn # type: ignore
-import numpy as np # type: ignore
-import umap # type: ignore
-from sklearn.preprocessing import MinMaxScaler # type: ignore
+from fastapi import FastAPI, Request, BackgroundTasks  # type: ignore
+from fastapi.responses import JSONResponse, StreamingResponse, FileResponse, RedirectResponse  # type: ignore
+from fastapi.middleware.cors import CORSMiddleware  # type: ignore
+import uvicorn  # type: ignore
+import numpy as np  # type: ignore
+import umap  # type: ignore
+from sklearn.preprocessing import MinMaxScaler  # type: ignore
 
 # Prometheus
-from prometheus_client import Summary # type: ignore
-from prometheus_fastapi_instrumentator import Instrumentator # type: ignore
-from prometheus_client import CollectorRegistry, Counter # type: ignore
+from prometheus_client import Summary  # type: ignore
+from prometheus_fastapi_instrumentator import Instrumentator  # type: ignore
+from prometheus_client import CollectorRegistry, Counter  # type: ignore
 
 from utils import (
     rag as Rag,
     tools as Tools,
-    Context, Conversation, Message, 
+    Context,
+    Conversation,
+    Message,
     Agent,
     Metrics,
     Tunables,
@@ -71,14 +76,25 @@ from utils import (
     logger,
 )
 
-CONTEXT_VERSION=2
+CONTEXT_VERSION = 2
 
 rags = [
-    { "name": "JPK", "enabled": True, "description": "Expert data about James Ketrenos, including work history, personal hobbies, and projects." },
+    {
+        "name": "JPK",
+        "enabled": True,
+        "description": "Expert data about James Ketrenos, including work history, personal hobbies, and projects.",
+    },
     # { "name": "LKML", "enabled": False, "description": "Full associative data for entire LKML mailing list archive." },
 ]
 
-REQUEST_TIME = Summary('request_processing_seconds', 'Time spent processing request')
+
+class QueryOptions(BaseModel):
+    prompt: str
+    tunables: Tunables = Field(default_factory=Tunables)
+    agent_options: Dict[str, Any] = Field(default={})
+
+
+REQUEST_TIME = Summary("request_processing_seconds", "Time spent processing request")
 
 system_message_old = f"""
 Launched on {datetime.now().isoformat()}.
@@ -107,6 +123,7 @@ You are provided with a <|resume|> which was generated by you, the <|context|> y
 Your task is to answer questions about the <|fact_check|> you generated based on the <|resume|> and <|context>.
 """
 
+
 def get_installed_ram():
     try:
         with open("/proc/meminfo", "r") as f:
@@ -117,21 +134,29 @@ def get_installed_ram():
     except Exception as e:
         return f"Error retrieving RAM: {e}"
 
+
 def get_graphics_cards():
     gpus = []
     try:
         # Run the ze-monitor utility
-        result = subprocess.run(["ze-monitor"], capture_output=True, text=True, check=True)
-        
+        result = subprocess.run(
+            ["ze-monitor"], capture_output=True, text=True, check=True
+        )
+
         # Clean up the output (remove leading/trailing whitespace and newlines)
         output = result.stdout.strip()
         for index in range(len(output.splitlines())):
-            result = subprocess.run(["ze-monitor", "--device", f"{index+1}", "--info"], capture_output=True, text=True, check=True)
+            result = subprocess.run(
+                ["ze-monitor", "--device", f"{index+1}", "--info"],
+                capture_output=True,
+                text=True,
+                check=True,
+            )
             gpu_info = result.stdout.strip().splitlines()
             gpu = {
                 "discrete": True,  # Assume it's discrete initially
                 "name": None,
-                "memory": None
+                "memory": None,
             }
             gpus.append(gpu)
             for line in gpu_info:
@@ -149,11 +174,12 @@ def get_graphics_cards():
                 if match:
                     gpu["discrete"] = False
                     continue
-        
+
         return gpus
     except Exception as e:
         return f"Error retrieving GPU info: {e}"
 
+
 def get_cpu_info():
     try:
         with open("/proc/cpuinfo", "r") as f:
@@ -165,6 +191,7 @@ def get_cpu_info():
     except Exception as e:
         return f"Error retrieving CPU info: {e}"
 
+
 def system_info(model):
     return {
         "System RAM": get_installed_ram(),
@@ -172,18 +199,19 @@ def system_info(model):
         "CPU": get_cpu_info(),
         "LLM Model": model,
         "Embedding Model": defines.embedding_model,
-        "Context length": defines.max_context
+        "Context length": defines.max_context,
     }
 
+
 # %%
 # Defaults
 OLLAMA_API_URL = defines.ollama_api_url
 MODEL_NAME = defines.model
-LOG_LEVEL="info"
-USE_TLS=False
-WEB_HOST="0.0.0.0"
-WEB_PORT=8911
-DEFAULT_HISTORY_LENGTH=5
+LOG_LEVEL = "info"
+USE_TLS = False
+WEB_HOST = "0.0.0.0"
+WEB_PORT = 8911
+DEFAULT_HISTORY_LENGTH = 5
 
 # %%
 # Globals
@@ -192,43 +220,65 @@ DEFAULT_HISTORY_LENGTH=5
 def create_system_message(prompt):
     return [{"role": "system", "content": prompt}]
 
+
 tool_log = []
 command_log = []
 model = None
 client = None
 web_server = None
 
+
 # %%
 # Cmd line overrides
 def parse_args():
     parser = argparse.ArgumentParser(description="AI is Really Cool")
-    parser.add_argument("--ollama-server", type=str, default=OLLAMA_API_URL, help=f"Ollama API endpoint. default={OLLAMA_API_URL}")
-    parser.add_argument("--ollama-model", type=str, default=MODEL_NAME, help=f"LLM model to use. default={MODEL_NAME}")
-    parser.add_argument("--web-host", type=str, default=WEB_HOST, help=f"Host to launch Flask web server. default={WEB_HOST} only if --web-disable not specified.")
-    parser.add_argument("--web-port", type=str, default=WEB_PORT, help=f"Port to launch Flask web server. default={WEB_PORT} only if --web-disable not specified.")
-    parser.add_argument("--level", type=str, choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], 
-                        default=LOG_LEVEL, help=f"Set the logging level. default={LOG_LEVEL}")
+    parser.add_argument(
+        "--ollama-server",
+        type=str,
+        default=OLLAMA_API_URL,
+        help=f"Ollama API endpoint. default={OLLAMA_API_URL}",
+    )
+    parser.add_argument(
+        "--ollama-model",
+        type=str,
+        default=MODEL_NAME,
+        help=f"LLM model to use. default={MODEL_NAME}",
+    )
+    parser.add_argument(
+        "--web-host",
+        type=str,
+        default=WEB_HOST,
+        help=f"Host to launch Flask web server. default={WEB_HOST} only if --web-disable not specified.",
+    )
+    parser.add_argument(
+        "--web-port",
+        type=str,
+        default=WEB_PORT,
+        help=f"Port to launch Flask web server. default={WEB_PORT} only if --web-disable not specified.",
+    )
+    parser.add_argument(
+        "--level",
+        type=str,
+        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
+        default=LOG_LEVEL,
+        help=f"Set the logging level. default={LOG_LEVEL}",
+    )
     return parser.parse_args()
 
 
-
 # %%
 
 
 # %%
 
+
 # %%
 def is_valid_uuid(value: str) -> bool:
     try:
         uuid_obj = uuid.UUID(value, version=4)
         return str(uuid_obj) == value
     except (ValueError, TypeError):
-        return False    
-
-
-
-
-
+        return False
 
 
 # %%
@@ -237,11 +287,13 @@ class WebServer:
     async def lifespan(self, app: FastAPI):
         # Start the file watcher
         self.observer, self.file_watcher = Rag.start_file_watcher(
-            llm=self.llm, 
+            llm=self.llm,
             watch_directory=defines.doc_dir,
-            recreate=False    # Don't recreate if exists
-        )                  
-        logger.info(f"API started with {self.file_watcher.collection.count()} documents in the collection")
+            recreate=False,  # Don't recreate if exists
+        )
+        logger.info(
+            f"API started with {self.file_watcher.collection.count()} documents in the collection"
+        )
         yield
         if self.observer:
             self.observer.stop()
@@ -256,13 +308,13 @@ class WebServer:
 
         # Keep the Instrumentator instance alive
         self.instrumentator = Instrumentator(registry=self.prometheus_collector)
-        
+
         # Instrument the FastAPI app
         self.instrumentator.instrument(self.app)
-        
+
         # Expose the /metrics endpoint
-        self.instrumentator.expose(self.app, endpoint="/metrics")        
-        
+        self.instrumentator.expose(self.app, endpoint="/metrics")
+
         self.contexts = {}
 
         self.llm = llm
@@ -271,15 +323,17 @@ class WebServer:
         self.file_watcher = None
         self.observer = None
 
-        self.ssl_enabled = os.path.exists(defines.key_path) and os.path.exists(defines.cert_path)
+        self.ssl_enabled = os.path.exists(defines.key_path) and os.path.exists(
+            defines.cert_path
+        )
 
         if self.ssl_enabled:
-            allow_origins=["https://battle-linux.ketrenos.com:3000"]
+            allow_origins = ["https://battle-linux.ketrenos.com:3000"]
         else:
-            allow_origins=["http://battle-linux.ketrenos.com:3000"]
+            allow_origins = ["http://battle-linux.ketrenos.com:3000"]
 
         logger.info(f"Allowed origins: {allow_origins}")
-            
+
         self.app.add_middleware(
             CORSMiddleware,
             allow_origins=allow_origins,
@@ -296,7 +350,7 @@ class WebServer:
             context = self.create_context()
             logger.info(f"Redirecting non-context to {context.id}")
             return RedirectResponse(url=f"/{context.id}", status_code=307)
-            #return JSONResponse({"redirect": f"/{context.id}"})
+            # return JSONResponse({"redirect": f"/{context.id}"})
 
         @self.app.put("/api/umap/{context_id}")
         async def put_umap(context_id: str, request: Request):
@@ -304,17 +358,21 @@ class WebServer:
             try:
                 if not self.file_watcher:
                     raise Exception("File watcher not initialized")
-            
+
                 context = self.upsert_context(context_id)
                 if not context:
-                    return JSONResponse({"error": f"Invalid context: {context_id}"}, status_code=400)
+                    return JSONResponse(
+                        {"error": f"Invalid context: {context_id}"}, status_code=400
+                    )
 
                 data = await request.json()
 
                 dimensions = data.get("dimensions", 2)
                 result = self.file_watcher.umap_collection
                 if not result:
-                    return JSONResponse({"error": "No UMAP collection found"}, status_code=404)
+                    return JSONResponse(
+                        {"error": "No UMAP collection found"}, status_code=404
+                    )
                 if dimensions == 2:
                     logger.info("Returning 2D UMAP")
                     umap_embedding = self.file_watcher.umap_embedding_2d
@@ -323,8 +381,10 @@ class WebServer:
                     umap_embedding = self.file_watcher.umap_embedding_3d
 
                 if len(umap_embedding) == 0:
-                    return JSONResponse({"error": "No UMAP embedding found"}, status_code=404)
-                
+                    return JSONResponse(
+                        {"error": "No UMAP embedding found"}, status_code=404
+                    )
+
                 result["embeddings"] = umap_embedding.tolist()
 
                 return JSONResponse(result)
@@ -339,43 +399,64 @@ class WebServer:
             logger.info(f"{request.method} {request.url.path}")
             if not self.file_watcher:
                 raise Exception("File watcher not initialized")
-        
+
             if not is_valid_uuid(context_id):
                 logger.warning(f"Invalid context_id: {context_id}")
                 return JSONResponse({"error": "Invalid context_id"}, status_code=400)
-            
+
             try:
                 data = await request.json()
                 query = data.get("query", "")
+                threshold = data.get("threshold", 0.5)
+                results = data.get("results", 10)
             except:
                 query = ""
+                threshold = 0.5
+                results = 10
             if not query:
-                return JSONResponse({"error": "No query provided for similarity search"}, status_code=400)
-
+                return JSONResponse(
+                    {"error": "No query provided for similarity search"},
+                    status_code=400,
+                )
             try:
-                chroma_results = self.file_watcher.find_similar(query=query, top_k=10)
+                chroma_results = self.file_watcher.find_similar(
+                    query=query, top_k=results, threshold=threshold
+                )
                 if not chroma_results:
                     return JSONResponse({"error": "No results found"}, status_code=404)
 
-                chroma_embedding = np.array(chroma_results["query_embedding"]).flatten()  # Ensure correct shape
+                chroma_embedding = np.array(
+                    chroma_results["query_embedding"]
+                ).flatten()  # Ensure correct shape
                 logger.info(f"Chroma embedding shape: {chroma_embedding.shape}")
 
-                umap_2d = self.file_watcher.umap_model_2d.transform([chroma_embedding])[0].tolist()
-                logger.info(f"UMAP 2D output: {umap_2d}, length: {len(umap_2d)}")  # Debug output
+                umap_2d = self.file_watcher.umap_model_2d.transform([chroma_embedding])[
+                    0
+                ].tolist()
+                logger.info(
+                    f"UMAP 2D output: {umap_2d}, length: {len(umap_2d)}"
+                )  # Debug output
 
-                umap_3d = self.file_watcher.umap_model_3d.transform([chroma_embedding])[0].tolist()
-                logger.info(f"UMAP 3D output: {umap_3d}, length: {len(umap_3d)}")  # Debug output
+                umap_3d = self.file_watcher.umap_model_3d.transform([chroma_embedding])[
+                    0
+                ].tolist()
+                logger.info(
+                    f"UMAP 3D output: {umap_3d}, length: {len(umap_3d)}"
+                )  # Debug output
+
+                return JSONResponse(
+                    {
+                        **chroma_results,
+                        "query": query,
+                        "umap_embedding_2d": umap_2d,
+                        "umap_embedding_3d": umap_3d,
+                    }
+                )
 
-                return JSONResponse({ 
-                    **chroma_results, 
-                    "query": query, 
-                    "umap_embedding_2d": umap_2d,
-                    "umap_embedding_3d": umap_3d
-                })
-            
             except Exception as e:
                 logger.error(e)
-                #return JSONResponse({"error": str(e)}, 500)
+                logging.error(traceback.format_exc())
+                return JSONResponse({"error": str(e)}, 500)
 
         @self.app.put("/api/reset/{context_id}/{agent_type}")
         async def put_reset(context_id: str, agent_type: str, request: Request):
@@ -386,7 +467,10 @@ class WebServer:
             context = self.upsert_context(context_id)
             agent = context.get_agent(agent_type)
             if not agent:
-                return JSONResponse({ "error": f"{agent_type} is not recognized", "context": context.id }, status_code=404)
+                return JSONResponse(
+                    {"error": f"{agent_type} is not recognized", "context": context.id},
+                    status_code=404,
+                )
 
             data = await request.json()
             try:
@@ -405,20 +489,29 @@ class WebServer:
                             response["tools"] = context.tools
                         case "history":
                             reset_map = {
-                                "job_description": ("job_description", "resume", "fact_check"),
+                                "job_description": (
+                                    "job_description",
+                                    "resume",
+                                    "fact_check",
+                                ),
                                 "resume": ("job_description", "resume", "fact_check"),
-                                "fact_check": ("job_description", "resume", "fact_check"),
+                                "fact_check": (
+                                    "job_description",
+                                    "resume",
+                                    "fact_check",
+                                ),
                                 "chat": ("chat",),
                             }
                             resets = reset_map.get(agent_type, ())
-
                             for mode in resets:
                                 tmp = context.get_agent(mode)
                                 if not tmp:
+                                    logger.info(
+                                        f"Agent {mode} not found for {context_id}"
+                                    )
                                     continue
                                 logger.info(f"Resetting {reset_operation} for {mode}")
-                                context.conversation = Conversation()
-                                context.context_tokens = round(len(str(agent.system_prompt)) * 3 / 4) # Estimate context usage
+                                tmp.conversation.reset()
                             response["history"] = []
                             response["context_used"] = agent.context_tokens
                         case "message_history_length":
@@ -427,71 +520,104 @@ class WebServer:
                             response["message_history_length"] = DEFAULT_HISTORY_LENGTH
 
                 if not response:
-                    return JSONResponse({ "error": "Usage: { reset: rags|tools|history|system_prompt}"})
+                    return JSONResponse(
+                        {"error": "Usage: { reset: rags|tools|history|system_prompt}"}
+                    )
                 else:
                     self.save_context(context_id)
                     return JSONResponse(response)
-        
-            except:
-                return JSONResponse({ "error": "Usage: { reset: rags|tools|history|system_prompt}"})
+
+            except Exception as e:
+                logger.error(f"Error in reset: {e}")
+                logger.error(traceback.format_exc())
+                return JSONResponse(
+                    {"error": "Usage: { reset: rags|tools|history|system_prompt}"}
+                )
 
         @self.app.put("/api/tunables/{context_id}")
         async def put_tunables(context_id: str, request: Request):
             logger.info(f"{request.method} {request.url.path}")
             try:
                 context = self.upsert_context(context_id)
-                
+
                 data = await request.json()
                 agent = context.get_agent("chat")
                 if not agent:
-                    return JSONResponse({ "error": f"chat is not recognized", "context": context.id }, status_code=404)
+                    return JSONResponse(
+                        {"error": f"chat is not recognized", "context": context.id},
+                        status_code=404,
+                    )
                 for k in data.keys():
                     match k:
                         case "tools":
                             # { "tools": [{ "tool": tool?.name, "enabled": tool.enabled }] }
-                            tools : list[dict[str, Any]] = data[k]
+                            tools: list[dict[str, Any]] = data[k]
                             if not tools:
-                                return JSONResponse({ "status": "error", "message": "Tools can not be empty." })
+                                return JSONResponse(
+                                    {
+                                        "status": "error",
+                                        "message": "Tools can not be empty.",
+                                    }
+                                )
                             for tool in tools:
                                 for context_tool in context.tools:
                                     if context_tool["function"]["name"] == tool["name"]:
                                         context_tool["enabled"] = tool["enabled"]
                             self.save_context(context_id)
-                            return JSONResponse({ "tools": [ { 
-                                **t["function"],
-                                "enabled": t["enabled"],
-                            } for t in context.tools] })
-                        
+                            return JSONResponse(
+                                {
+                                    "tools": [
+                                        {
+                                            **t["function"],
+                                            "enabled": t["enabled"],
+                                        }
+                                        for t in context.tools
+                                    ]
+                                }
+                            )
+
                         case "rags":
                             # { "rags": [{ "tool": tool?.name, "enabled": tool.enabled }] }
-                            rags : list[dict[str, Any]] = data[k]
+                            rags: list[dict[str, Any]] = data[k]
                             if not rags:
-                                return JSONResponse({ "status": "error", "message": "RAGs can not be empty." })
+                                return JSONResponse(
+                                    {
+                                        "status": "error",
+                                        "message": "RAGs can not be empty.",
+                                    }
+                                )
                             for rag in rags:
                                 for context_rag in context.rags:
                                     if context_rag["name"] == rag["name"]:
                                         context_rag["enabled"] = rag["enabled"]
                             self.save_context(context_id)
-                            return JSONResponse({ "rags": context.rags })
-                            
+                            return JSONResponse({"rags": context.rags})
+
                         case "system_prompt":
                             system_prompt = data[k].strip()
                             if not system_prompt:
-                                return JSONResponse({ "status": "error", "message": "System prompt can not be empty." })
-                            agent.system_prompt =  system_prompt
+                                return JSONResponse(
+                                    {
+                                        "status": "error",
+                                        "message": "System prompt can not be empty.",
+                                    }
+                                )
+                            agent.system_prompt = system_prompt
                             self.save_context(context_id)
-                            return JSONResponse({ "system_prompt": system_prompt })
+                            return JSONResponse({"system_prompt": system_prompt})
                         case "message_history_length":
                             value = max(0, int(data[k]))
                             context.message_history_length = value
                             self.save_context(context_id)
-                            return JSONResponse({ "message_history_length": value })
+                            return JSONResponse({"message_history_length": value})
                         case _:
-                            return JSONResponse({ "error": f"Unrecognized tunable {k}"}, status_code=404)
+                            return JSONResponse(
+                                {"error": f"Unrecognized tunable {k}"}, status_code=404
+                            )
             except Exception as e:
                 logger.error(f"Error in put_tunables: {e}")
-                return JSONResponse({"error": str(e)}, status_code=500)        
-        
+                return JSONResponse({"error": str(e)}, status_code=500)
+
         @self.app.get("/api/tunables/{context_id}")
         async def get_tunables(context_id: str, request: Request):
             logger.info(f"{request.method} {request.url.path}")
@@ -501,85 +627,94 @@ class WebServer:
             context = self.upsert_context(context_id)
             agent = context.get_agent("chat")
             if not agent:
-                return JSONResponse({ "error": f"chat is not recognized", "context": context.id }, status_code=404)
-            return JSONResponse({ 
-                "system_prompt": agent.system_prompt,
-                "message_history_length": context.message_history_length,
-                "rags": context.rags,
-                "tools": [ { 
-                    **t["function"],
-                    "enabled": t["enabled"],
-                } for t in context.tools ]
-            })
+                return JSONResponse(
+                    {"error": f"chat is not recognized", "context": context.id},
+                    status_code=404,
+                )
+            return JSONResponse(
+                {
+                    "system_prompt": agent.system_prompt,
+                    "message_history_length": context.message_history_length,
+                    "rags": context.rags,
+                    "tools": [
+                        {
+                            **t["function"],
+                            "enabled": t["enabled"],
+                        }
+                        for t in context.tools
+                    ],
+                }
+            )
 
         @self.app.get("/api/system-info/{context_id}")
         async def get_system_info(context_id: str, request: Request):
             logger.info(f"{request.method} {request.url.path}")
             return JSONResponse(system_info(self.model))
 
-        @self.app.post("/api/chat/{context_id}/{agent_type}")
-        async def post_chat_endpoint(context_id: str, agent_type: str, request: Request):
+        @self.app.post("/api/{agent_type}/{context_id}")
+        async def post_agent_endpoint(
+            agent_type: str, context_id: str, request: Request
+        ):
             logger.info(f"{request.method} {request.url.path}")
-            if not is_valid_uuid(context_id):
-                logger.warning(f"Invalid context_id: {context_id}")
-                return JSONResponse({"error": "Invalid context_id"}, status_code=400)
 
             try:
                 context = self.upsert_context(context_id)
-                agent = context.get_agent(agent_type)
             except Exception as e:
-                logger.info(f"Attempt to create agent type: {agent_type} failed", e)
-                return JSONResponse({"error": f"{agent_type} is not recognized or context {context_id} is invalid "}, status_code=404)
+                error = {
+                    "error": f"Unable to create or access context {context_id}: {e}"
+                }
+                logger.info(error)
+                return JSONResponse(error, status_code=404)
 
             try:
-                query = await request.json()
-                prompt = query["prompt"]
-                if not isinstance(prompt, str) or len(prompt) == 0:
-                    logger.info(f"Prompt is empty")
-                    return JSONResponse({"error": "Prompt cannot be empty"}, status_code=400)
+                data = await request.json()
+                query: QueryOptions = QueryOptions(**data)
             except Exception as e:
-                logger.info(f"Attempt to parse request: {str(e)}.")
-                return JSONResponse({"error": f"Attempt to parse request: {str(e)}."}, status_code=400)
-                
-            try:
-                options = Tunables(**query["options"]) if "options" in query else None
-            except Exception as e:
-                logger.info(f"Attempt to set tunables failed: {query['options']}.", e)
-                return JSONResponse({"error": f"Invalid options: {query['options']}"}, status_code=400)
+                error = {"error": f"Attempt to parse request: {e}"}
+                logger.info(error)
+                return JSONResponse(error, status_code=400)
 
-            if not agent:
-                match agent_type:
-                    case "job_description":
-                        logger.info(f"Agent {agent_type} not found. Returning empty history.")
-                        agent = context.get_or_create_agent("job_description", job_description=prompt)
-                    case _:
-                        logger.info(f"Invalid agent creation sequence for {agent_type}. Returning error.")
-                        return JSONResponse({"error": f"{agent_type} is not recognized", "context": context.id}, status_code=404)
             try:
+                agent = context.get_or_create_agent(agent_type, **query.agent_options)
+            except Exception as e:
+                error = {
+                    "error": f"Attempt to create agent type: {agent_type} failed: {e}"
+                }
+                return JSONResponse(error, status_code=404)
+
+            try:
+
                 async def flush_generator():
                     logger.info(f"{agent.agent_type} - {inspect.stack()[0].function}")
                     try:
                         start_time = time.perf_counter()
-                        async for message in self.generate_response(context=context, agent=agent, prompt=prompt, options=options):
-                            if message.status != "done":
+                        async for message in self.generate_response(
+                            context=context,
+                            agent=agent,
+                            prompt=query.prompt,
+                            options=query.options,
+                        ):
+                            if message.status != "done" and message.status != "partial":
                                 if message.status == "streaming":
                                     result = {
                                         "status": "streaming",
                                         "chunk": message.chunk,
-                                        "remaining_time": LLM_TIMEOUT - (time.perf_counter() - start_time)
+                                        "remaining_time": LLM_TIMEOUT
+                                        - (time.perf_counter() - start_time),
                                     }
                                 else:
                                     start_time = time.perf_counter()
                                     result = {
                                         "status": message.status,
                                         "response": message.response,
-                                        "remaining_time": LLM_TIMEOUT
+                                        "remaining_time": LLM_TIMEOUT,
                                     }
                             else:
-                                logger.info(f"Message complete. Providing full response.")
+                                logger.info(f"Providing {message.status} response.")
                                 try:
-                                    message.response = message.response
-                                    result = message.model_dump(by_alias=True, mode='json')
+                                    result = message.model_dump(
+                                        by_alias=True, mode="json"
+                                    )
                                 except Exception as e:
                                     result = {"status": "error", "response": str(e)}
                                     yield json.dumps(result) + "\n"
@@ -589,26 +724,27 @@ class WebServer:
                             result = json.dumps(result) + "\n"
                             message.network_packets += 1
                             message.network_bytes += len(result)
+                            yield result
                             if await request.is_disconnected():
                                 logger.info("Disconnect detected. Aborting generation.")
                                 context.processing = False
                                 # Save context on completion or error
-                                message.prompt = prompt
+                                message.prompt = query.prompt
                                 message.status = "error"
-                                message.response = "Client disconnected during generation."
+                                message.response = (
+                                    "Client disconnected during generation."
+                                )
                                 agent.conversation.add(message)
                                 self.save_context(context_id)
                                 return
 
-                            yield result
-
                             current_time = time.perf_counter()
                             if current_time - start_time > LLM_TIMEOUT:
                                 message.status = "error"
                                 message.response = f"Processing time ({LLM_TIMEOUT}s) exceeded for single LLM inference (likely due to LLM getting stuck.) You will need to retry your query."
                                 message.partial_response = message.response
                                 logger.info(message.response + " Ending session")
-                                result = message.model_dump(by_alias=True, mode='json')
+                                result = message.model_dump(by_alias=True, mode="json")
                                 result = json.dumps(result) + "\n"
                                 yield result
 
@@ -620,13 +756,13 @@ class WebServer:
                             await asyncio.sleep(0)
                     except Exception as e:
                         context.processing = False
-                        logger.error(f"Error in process_generator: {e}")
+                        logger.error(f"Error in generate_response: {e}")
                         logger.error(traceback.format_exc())
                         yield json.dumps({"status": "error", "response": str(e)}) + "\n"
                     finally:
                         # Save context on completion or error
                         self.save_context(context_id)
-                                    
+
                 # Return StreamingResponse with appropriate headers
                 return StreamingResponse(
                     flush_generator(),
@@ -634,8 +770,8 @@ class WebServer:
                     headers={
                         "Cache-Control": "no-cache",
                         "Connection": "keep-alive",
-                        "X-Accel-Buffering": "no"  # Prevents Nginx buffering if you're using it
-                    }
+                        "X-Accel-Buffering": "no",  # Prevents Nginx buffering if you're using it
+                    },
                 )
             except Exception as e:
                 context.processing = False
@@ -647,7 +783,7 @@ class WebServer:
             try:
                 context = self.create_context()
                 logger.info(f"Generated new agent as {context.id}")
-                return JSONResponse({ "id": context.id })
+                return JSONResponse({"id": context.id})
             except Exception as e:
                 logger.error(f"get_history error: {str(e)}")
                 logger.error(traceback.format_exc())
@@ -660,16 +796,21 @@ class WebServer:
                 context = self.upsert_context(context_id)
                 agent = context.get_agent(agent_type)
                 if not agent:
-                    logger.info(f"Agent {agent_type} not found. Returning empty history.")
-                    return JSONResponse({ "messages": [] })
-                logger.info(f"History for {agent_type} contains {len(agent.conversation)} entries.")
+                    logger.info(
+                        f"Agent {agent_type} not found. Returning empty history."
+                    )
+                    return JSONResponse({"messages": []})
+                logger.info(
+                    f"History for {agent_type} contains {len(agent.conversation)} entries."
+                )
                 return agent.conversation
             except Exception as e:
                 logger.error(f"get_history error: {str(e)}")
                 import traceback
+
                 logger.error(traceback.format_exc())
                 return JSONResponse({"error": str(e)}, status_code=404)
-            
+
         @self.app.get("/api/tools/{context_id}")
         async def get_tools(context_id: str, request: Request):
             logger.info(f"{request.method} {request.url.path}")
@@ -692,10 +833,11 @@ class WebServer:
                         tool["enabled"] = enabled
                         self.save_context(context_id)
                         return JSONResponse(context.tools)
-                return JSONResponse({ "status": f"{modify} not found in tools." }, status_code=404)
+                return JSONResponse(
+                    {"status": f"{modify} not found in tools."}, status_code=404
+                )
             except:
-                return JSONResponse({ "status": "error" }, 405)
-
+                return JSONResponse({"status": "error"}, 405)
 
         @self.app.get("/api/context-status/{context_id}/{agent_type}")
         async def get_context_status(context_id, agent_type: str, request: Request):
@@ -706,8 +848,15 @@ class WebServer:
             context = self.upsert_context(context_id)
             agent = context.get_agent(agent_type)
             if not agent:
-                return JSONResponse({"context_used": 0, "max_context": defines.max_context})
-            return JSONResponse({"context_used": agent.context_tokens, "max_context": defines.max_context})
+                return JSONResponse(
+                    {"context_used": 0, "max_context": defines.max_context}
+                )
+            return JSONResponse(
+                {
+                    "context_used": agent.context_tokens,
+                    "max_context": defines.max_context,
+                }
+            )
 
         @self.app.get("/api/health")
         async def health_check():
@@ -719,18 +868,18 @@ class WebServer:
             if os.path.exists(full_path) and os.path.isfile(full_path):
                 logger.info(f"Serve static request for {full_path}")
                 return FileResponse(full_path)
-            
+
             logger.info(f"Serve index.html for {path}")
             return FileResponse(os.path.join(defines.static_content, "index.html"))
-        
+
     def save_context(self, context_id):
         """
         Serialize a Python dictionary to a file in the agents directory.
-        
+
         Args:
             data: Dictionary containing the agent data
             context_id: UUID string for the context. If it doesn't exist, it is created
-            
+
         Returns:
             The context_id used for the file
         """
@@ -739,16 +888,16 @@ class WebServer:
         # Create agents directory if it doesn't exist
         if not os.path.exists(defines.context_dir):
             os.makedirs(defines.context_dir)
-        
+
         # Create the full file path
         file_path = os.path.join(defines.context_dir, context_id)
-        
+
         # Serialize the data to JSON and write to file
         with open(file_path, "w") as f:
             f.write(context.model_dump_json(by_alias=True))
 
         return context_id
-    
+
     def load_or_create_context(self, context_id) -> Context:
         """
         Load a context from a file in the context directory or create a new one if it doesn't exist.
@@ -761,7 +910,7 @@ class WebServer:
             raise Exception("File watcher not initialized")
 
         file_path = os.path.join(defines.context_dir, context_id)
-        
+
         # Check if the file exists
         if not os.path.exists(file_path):
             logger.info(f"Context file {file_path} not found. Creating new context.")
@@ -770,13 +919,16 @@ class WebServer:
             # Read and deserialize the data
             with open(file_path, "r") as f:
                 content = f.read()
-                logger.info(f"Loading context from {file_path}, content length: {len(content)}")
+                logger.info(
+                    f"Loading context from {file_path}, content length: {len(content)}"
+                )
                 import json
+
                 try:
                     # Try parsing as JSON first to ensure valid JSON
                     json_data = json.loads(content)
                     logger.info("JSON parsed successfully, attempting model validation")
-                    
+
                     # Validate from JSON (no prometheus_collector or file_watcher)
                     context = Context.model_validate(json_data)
 
@@ -787,7 +939,9 @@ class WebServer:
                     # Now set context on agents manually
                     agent_types = [agent.agent_type for agent in context.agents]
                     if len(agent_types) != len(set(agent_types)):
-                        raise ValueError("Context cannot contain multiple agents of the same agent_type")
+                        raise ValueError(
+                            "Context cannot contain multiple agents of the same agent_type"
+                        )
                     for agent in context.agents:
                         agent.set_context(context)
 
@@ -799,13 +953,18 @@ class WebServer:
                 except Exception as e:
                     logger.error(f"Error validating context: {str(e)}")
                     import traceback
+
                     logger.error(traceback.format_exc())
                     # Fallback to creating a new context
-                    self.contexts[context_id] = Context(id=context_id, file_watcher=self.file_watcher, prometheus_collector=self.prometheus_collector)
-            
+                    self.contexts[context_id] = Context(
+                        id=context_id,
+                        file_watcher=self.file_watcher,
+                        prometheus_collector=self.prometheus_collector,
+                    )
+
         return self.contexts[context_id]
 
-    def create_context(self, context_id = None) -> Context:
+    def create_context(self, context_id=None) -> Context:
         """
         Create a new context with a unique ID and default settings.
         Args:
@@ -818,8 +977,12 @@ class WebServer:
         if not context_id:
             context_id = str(uuid4())
         logger.info(f"Creating new context with ID: {context_id}")
-        context = Context(id=context_id, file_watcher=self.file_watcher, prometheus_collector=self.prometheus_collector)
-        
+        context = Context(
+            id=context_id,
+            file_watcher=self.file_watcher,
+            prometheus_collector=self.prometheus_collector,
+        )
+
         if os.path.exists(defines.resume_doc):
             context.user_resume = open(defines.resume_doc, "r").read()
         context.get_or_create_agent(agent_type="chat")
@@ -835,7 +998,7 @@ class WebServer:
         self.save_context(context.id)
         return context
 
-    def upsert_context(self, context_id = None) -> Context:
+    def upsert_context(self, context_id=None) -> Context:
         """
         Upsert a context based on the provided context_id.
         Args:
@@ -847,18 +1010,20 @@ class WebServer:
         if not context_id:
             logger.warning("No context ID provided. Creating a new context.")
             return self.create_context()
-        
+
         if context_id in self.contexts:
             return self.contexts[context_id]
-        
+
         logger.info(f"Context {context_id} is not yet loaded.")
         return self.load_or_create_context(context_id)
-        
+
     @REQUEST_TIME.time()
-    async def generate_response(self, context : Context, agent : Agent, prompt : str, options: Tunables | None) -> AsyncGenerator[Message, None]:
+    async def generate_response(
+        self, context: Context, agent: Agent, prompt: str, options: Tunables | None
+    ) -> AsyncGenerator[Message, None]:
         if not self.file_watcher:
             raise Exception("File watcher not initialized")
-        
+
         agent_type = agent.get_agent_type()
         logger.info(f"generate_response: type - {agent_type}")
         message = Message(prompt=prompt, options=agent.tunables)
@@ -876,11 +1041,13 @@ class WebServer:
         async for message in agent.process_message(self.llm, self.model, message):
             if message.status != "done":
                 yield message
-    
+
         if message.status == "error":
             return
 
-        logger.info(f"{agent_type}.process_message: {message.status} {f'...{message.response[-20:]}' if len(message.response) > 20 else message.response}")
+        logger.info(
+            f"{agent_type}.process_message: {message.status} {f'...{message.response[-20:]}' if len(message.response) > 20 else message.response}"
+        )
         message.status = "done"
         yield message
         return
@@ -894,25 +1061,22 @@ class WebServer:
                     host=host,
                     port=port,
                     log_config=None,
-                    ssl_keyfile=defines.key_path, 
-                    ssl_certfile=defines.cert_path
+                    ssl_keyfile=defines.key_path,
+                    ssl_certfile=defines.cert_path,
                 )
             else:
                 logger.info(f"Starting web server at http://{host}:{port}")
-                uvicorn.run(
-                    self.app,
-                    host=host,
-                    port=port,
-                    log_config=None
-                )
+                uvicorn.run(self.app, host=host, port=port, log_config=None)
         except KeyboardInterrupt:
             if self.observer:
                 self.observer.stop()
         if self.observer:
             self.observer.join()
 
+
 # %%
 
+
 # Main function to run everything
 def main():
     global model
@@ -923,23 +1087,16 @@ def main():
     # Setup logging based on the provided level
     logger.setLevel(args.level.upper())
 
-    warnings.filterwarnings(
-        "ignore",
-        category=FutureWarning,
-        module="sklearn.*"
-    )
+    warnings.filterwarnings("ignore", category=FutureWarning, module="sklearn.*")
 
-    warnings.filterwarnings(
-        "ignore",
-        category=UserWarning,
-        module="umap.*"
-    )
+    warnings.filterwarnings("ignore", category=UserWarning, module="umap.*")
 
-    llm = ollama.Client(host=args.ollama_server) # type: ignore
+    llm = ollama.Client(host=args.ollama_server)  # type: ignore
     model = args.ollama_model
 
     web_server = WebServer(llm, model)
-    
+
     web_server.run(host=args.web_host, port=args.web_port, use_reloader=False)
 
+
 main()
diff --git a/src/tests/test-agent.py b/src/tests/test-agent.py
index 3f4386d..0c51907 100644
--- a/src/tests/test-agent.py
+++ b/src/tests/test-agent.py
@@ -1,29 +1,25 @@
-from .. utils import logger
+from ..utils import logger
 
 import ollama
 
-from .. utils import (
-  rag as Rag,
-  Context,
-  defines
-)
+from ..utils import rag as Rag, Context, defines
 
 import json
 
 llm = ollama.Client(host=defines.ollama_api_url)
-    
+
 observer, file_watcher = Rag.start_file_watcher(
-    llm=llm, 
-    watch_directory=defines.doc_dir,
-    recreate=False    # Don't recreate if exists
+    llm=llm, watch_directory=defines.doc_dir, recreate=False  # Don't recreate if exists
 )
 
 context = Context(file_watcher=file_watcher)
-data = context.model_dump(mode='json')
+data = context.model_dump(mode="json")
 context = Context.model_validate_json(json.dumps(data))
 context.file_watcher = file_watcher
 
-agent = context.get_or_create_agent("chat", system_prompt="You are a helpful assistant.")
+agent = context.get_or_create_agent(
+    "chat", system_prompt="You are a helpful assistant."
+)
 # logger.info(f"data: {data}")
 # logger.info(f"agent: {agent}")
 agent_type = agent.get_agent_type()
@@ -32,7 +28,7 @@ logger.info(f"system_prompt: {agent.system_prompt}")
 
 agent.system_prompt = "Eat more tomatoes."
 
-data = context.model_dump(mode='json')
+data = context.model_dump(mode="json")
 context = Context.model_validate_json(json.dumps(data))
 context.file_watcher = file_watcher
 
diff --git a/src/tests/test-context-routing.py b/src/tests/test-context-routing.py
index b8adab4..7ed01e2 100644
--- a/src/tests/test-context-routing.py
+++ b/src/tests/test-context-routing.py
@@ -1,109 +1,139 @@
-from fastapi import FastAPI, Request, Depends, Query # type: ignore
-from fastapi.responses import RedirectResponse, JSONResponse # type: ignore
+from fastapi import FastAPI, Request, Depends, Query  # type: ignore
+from fastapi.responses import RedirectResponse, JSONResponse  # type: ignore
 from uuid import UUID, uuid4
 import logging
 import traceback
 from typing import Callable, Optional
-from anyio.to_thread import run_sync # type: ignore
+from anyio.to_thread import run_sync  # type: ignore
 
 logger = logging.getLogger(__name__)
 
+
 class RedirectToContext(Exception):
     def __init__(self, url: str):
         self.url = url
         logger.info(f"Redirect to Context: {url}")
         super().__init__(f"Redirect to Context: {url}")
 
+
 class ContextRouteManager:
     def __init__(self, app: FastAPI):
         self.app = app
         self.setup_handlers()
-    
+
     def setup_handlers(self):
         @self.app.exception_handler(RedirectToContext)
         async def handle_context_redirect(request: Request, exc: RedirectToContext):
             logger.info(f"Handling redirect to {exc.url}")
             return RedirectResponse(url=exc.url, status_code=307)
-    
-    def ensure_context(self, route_name: str = "context_id") -> Callable[[Request], Optional[UUID]]:
+
+    def ensure_context(
+        self, route_name: str = "context_id"
+    ) -> Callable[[Request], Optional[UUID]]:
         logger.info(f"Setting up context dependency for route parameter: {route_name}")
-        
+
         async def _ensure_context_dependency(request: Request) -> Optional[UUID]:
-            logger.info(f"Entering ensure_context_dependency, Request URL: {request.url}")
+            logger.info(
+                f"Entering ensure_context_dependency, Request URL: {request.url}"
+            )
             logger.info(f"Path params: {request.path_params}")
-            
+
             path_params = request.path_params
             route_value = path_params.get(route_name)
             logger.info(f"route_value: {route_value!r}, type: {type(route_value)}")
-            
-            if route_value is None or not isinstance(route_value, str) or not route_value.strip():
+
+            if (
+                route_value is None
+                or not isinstance(route_value, str)
+                or not route_value.strip()
+            ):
                 logger.info(f"route_value is invalid, generating new UUID")
-                path = request.url.path.rstrip('/')
+                path = request.url.path.rstrip("/")
                 new_context = await run_sync(uuid4)
                 redirect_url = f"{path}/{new_context}"
                 logger.info(f"Redirecting to {redirect_url}")
                 raise RedirectToContext(redirect_url)
-            
+
             logger.info(f"Attempting to parse route_value as UUID: {route_value}")
             try:
                 route_context = await run_sync(UUID, route_value)
                 logger.info(f"Successfully parsed UUID: {route_context}")
                 return route_context
             except ValueError as e:
-                logger.error(f"Failed to parse UUID from route_value: {route_value!r}, error: {str(e)}")
-                path = request.url.path.rstrip('/')
+                logger.error(
+                    f"Failed to parse UUID from route_value: {route_value!r}, error: {str(e)}"
+                )
+                path = request.url.path.rstrip("/")
                 new_context = await run_sync(uuid4)
                 redirect_url = f"{path}/{new_context}"
                 logger.info(f"Invalid UUID, redirecting to {redirect_url}")
                 raise RedirectToContext(redirect_url)
 
-        return _ensure_context_dependency # type: ignore
-    
+        return _ensure_context_dependency  # type: ignore
+
     def route_pattern(self, path: str, *dependencies, **kwargs):
         logger.info(f"Registering route: {path}")
         ensure_context = self.ensure_context()
-        
+
         def decorator(func):
             all_dependencies = list(dependencies)
             all_dependencies.append(Depends(ensure_context))
-            logger.info(f"Route {path} registered with dependencies: {all_dependencies}")
+            logger.info(
+                f"Route {path} registered with dependencies: {all_dependencies}"
+            )
             return self.app.get(path, dependencies=all_dependencies, **kwargs)(func)
-        
+
         return decorator
 
+
 app = FastAPI(redirect_slashes=True)
 
+
 @app.exception_handler(Exception)
 async def global_exception_handler(request: Request, exc: Exception):
     logger.error(f"Unhandled exception: {str(exc)}")
     logger.error(f"Request URL: {request.url}, Path params: {request.path_params}")
     logger.error(f"Stack trace: {''.join(traceback.format_tb(exc.__traceback__))}")
     return JSONResponse(
-        status_code=500,
-        content={"error": "Internal server error", "detail": str(exc)}
+        status_code=500, content={"error": "Internal server error", "detail": str(exc)}
     )
 
+
 @app.middleware("http")
 async def log_requests(request: Request, call_next):
-    logger.info(f"Incoming request: {request.method} {request.url}, Path params: {request.path_params}")
+    logger.info(
+        f"Incoming request: {request.method} {request.url}, Path params: {request.path_params}"
+    )
     response = await call_next(request)
     return response
 
+
 context_router = ContextRouteManager(app)
 
+
 @context_router.route_pattern("/api/history/{context_id}")
-async def get_history(request: Request, context_id: UUID = Depends(context_router.ensure_context()), agent_type: str = Query(..., description="Type of agent to retrieve history for")):
+async def get_history(
+    request: Request,
+    context_id: UUID = Depends(context_router.ensure_context()),
+    agent_type: str = Query(..., description="Type of agent to retrieve history for"),
+):
     logger.info(f"{request.method} {request.url.path} with context_id: {context_id}")
     return {"context_id": str(context_id), "agent_type": agent_type}
 
+
 @app.get("/api/history")
-async def redirect_history(request: Request, agent_type: str = Query(..., description="Type of agent to retrieve history for")):
-    path = request.url.path.rstrip('/')
+async def redirect_history(
+    request: Request,
+    agent_type: str = Query(..., description="Type of agent to retrieve history for"),
+):
+    path = request.url.path.rstrip("/")
     new_context = uuid4()
     redirect_url = f"{path}/{new_context}?agent_type={agent_type}"
     logger.info(f"Redirecting from /api/history to {redirect_url}")
     return RedirectResponse(url=redirect_url, status_code=307)
 
+
 if __name__ == "__main__":
-    import uvicorn # type: ignore
-    uvicorn.run(app, host="0.0.0.0", port=8900)
\ No newline at end of file
+    import uvicorn  # type: ignore
+
+    uvicorn.run(app, host="0.0.0.0", port=8900)
diff --git a/src/tests/test-context.py b/src/tests/test-context.py
index ef90bba..150bff2 100644
--- a/src/tests/test-context.py
+++ b/src/tests/test-context.py
@@ -1,28 +1,24 @@
 # From /opt/backstory run:
 # python -m src.tests.test-context
 import os
+
 os.environ["TORCH_CPP_LOG_LEVEL"] = "ERROR"
 import warnings
+
 warnings.filterwarnings("ignore", message="Couldn't find ffmpeg or avconv")
 
 import ollama
 
-from .. utils import (
-  rag as Rag,
-  Context,
-  defines
-)
+from ..utils import rag as Rag, Context, defines
 
 import json
 
-llm = ollama.Client(host=defines.ollama_api_url) # type: ignore
-    
+llm = ollama.Client(host=defines.ollama_api_url)  # type: ignore
+
 observer, file_watcher = Rag.start_file_watcher(
-    llm=llm, 
-    watch_directory=defines.doc_dir,
-    recreate=False    # Don't recreate if exists
+    llm=llm, watch_directory=defines.doc_dir, recreate=False  # Don't recreate if exists
 )
 
 context = Context(file_watcher=file_watcher)
-data = context.model_dump(mode='json')
+data = context.model_dump(mode="json")
 context = Context.from_json(json.dumps(data), file_watcher=file_watcher)
diff --git a/src/tests/test-message.py b/src/tests/test-message.py
index e8bec94..6de3b5b 100644
--- a/src/tests/test-message.py
+++ b/src/tests/test-message.py
@@ -1,13 +1,11 @@
 # From /opt/backstory run:
 # python -m src.tests.test-message
-from .. utils import logger
+from ..utils import logger
 
-from .. utils import (
-  Message
-)
+from ..utils import Message
 
 import json
 
 prompt = "This is a test"
 message = Message(prompt=prompt)
-print(message.model_dump(mode='json'))
+print(message.model_dump(mode="json"))
diff --git a/src/tests/test-metrics.py b/src/tests/test-metrics.py
index 1bb0421..d989167 100644
--- a/src/tests/test-metrics.py
+++ b/src/tests/test-metrics.py
@@ -1,8 +1,6 @@
 # From /opt/backstory run:
 # python -m src.tests.test-metrics
-from .. utils import (
-  Metrics
-)
+from ..utils import Metrics
 
 import json
 
@@ -13,7 +11,7 @@ metrics = Metrics()
 metrics.prepare_count.labels(agent="chat").inc()
 metrics.prepare_duration.labels(agent="prepare").observe(0.45)
 
-json = metrics.model_dump(mode='json')
+json = metrics.model_dump(mode="json")
 metrics = Metrics.model_validate(json)
 
-print(metrics)
\ No newline at end of file
+print(metrics)
diff --git a/src/utils/__init__.py b/src/utils/__init__.py
index 0cbb169..bcf65e4 100644
--- a/src/utils/__init__.py
+++ b/src/utils/__init__.py
@@ -1,44 +1,51 @@
 from __future__ import annotations
-from pydantic import BaseModel # type: ignore
+from pydantic import BaseModel  # type: ignore
 import importlib
 
-from . import defines  
-from . context import Context
-from . conversation import Conversation
-from . message import Message, Tunables
-from . rag import ChromaDBFileWatcher, start_file_watcher
-from . setup_logging import setup_logging
-from . agents import class_registry, AnyAgent, Agent, __all__ as agents_all
-from . metrics import Metrics
+from . import defines
+from .context import Context
+from .conversation import Conversation
+from .message import Message, Tunables
+from .rag import ChromaDBFileWatcher, start_file_watcher
+from .setup_logging import setup_logging
+from .agents import class_registry, AnyAgent, Agent, __all__ as agents_all
+from .metrics import Metrics
 
 __all__ = [
-  'Agent', 
-  'Tunables',
-  'Context', 
-  'Conversation', 
-  'Message', 
-  'Metrics',
-  'ChromaDBFileWatcher', 
-  'start_file_watcher',
-  'logger',
+    "Agent",
+    "Tunables",
+    "Context",
+    "Conversation",
+    "Message",
+    "Metrics",
+    "ChromaDBFileWatcher",
+    "start_file_watcher",
+    "logger",
 ]
 
-__all__.extend(agents_all) # type: ignore
+__all__.extend(agents_all)  # type: ignore
 
 logger = setup_logging(level=defines.logging_level)
 
+
 def rebuild_models():
     for class_name, (module_name, _) in class_registry.items():
         try:
             module = importlib.import_module(module_name)
             cls = getattr(module, class_name, None)
-            
+
             logger.debug(f"Checking: {class_name} in module {module_name}")
             logger.debug(f"  cls: {True if cls else False}")
             logger.debug(f"  isinstance(cls, type): {isinstance(cls, type)}")
-            logger.debug(f"  issubclass(cls, BaseModel): {issubclass(cls, BaseModel) if cls else False}")
-            logger.debug(f"  issubclass(cls, AnyAgent): {issubclass(cls, AnyAgent) if cls else False}")
-            logger.debug(f"  cls is not AnyAgent: {cls is not AnyAgent if cls else True}")
+            logger.debug(
+                f"  issubclass(cls, BaseModel): {issubclass(cls, BaseModel) if cls else False}"
+            )
+            logger.debug(
+                f"  issubclass(cls, AnyAgent): {issubclass(cls, AnyAgent) if cls else False}"
+            )
+            logger.debug(
+                f"  cls is not AnyAgent: {cls is not AnyAgent if cls else True}"
+            )
 
             if (
                 cls
@@ -48,13 +55,15 @@ def rebuild_models():
                 and cls is not AnyAgent
             ):
                 logger.debug(f"Rebuilding {class_name} from {module_name}")
-                from . agents import Agent
-                from . context import Context
+                from .agents import Agent
+                from .context import Context
+
                 cls.model_rebuild()
         except ImportError as e:
             logger.error(f"Failed to import module {module_name}: {e}")
         except Exception as e:
             logger.error(f"Error processing {class_name} in {module_name}: {e}")
 
+
 # Call this after all modules are imported
 rebuild_models()
diff --git a/src/utils/agents/__init__.py b/src/utils/agents/__init__.py
index 99c7ba8..8a3225c 100644
--- a/src/utils/agents/__init__.py
+++ b/src/utils/agents/__init__.py
@@ -4,19 +4,21 @@ import importlib
 import pathlib
 import inspect
 
-from . types import agent_registry
-from .. setup_logging import setup_logging
+from .types import agent_registry
+from ..setup_logging import setup_logging
 from .. import defines
-from . base import Agent
+from .base import Agent
 
 logger = setup_logging(defines.logging_level)
 
-__all__ = [ "AnyAgent", "Agent", "agent_registry", "class_registry" ]
+__all__ = ["AnyAgent", "Agent", "agent_registry", "class_registry"]
 
 # Type alias for Agent or any subclass
 AnyAgent: TypeAlias = Agent  # BaseModel covers Agent and subclasses
 
-class_registry: Dict[str, Tuple[str, str]] = {}  # Maps class_name to (module_name, class_name)
+class_registry: Dict[str, Tuple[str, str]] = (
+    {}
+)  # Maps class_name to (module_name, class_name)
 
 package_dir = pathlib.Path(__file__).parent
 package_name = __name__
@@ -42,7 +44,7 @@ for path in package_dir.glob("*.py"):
                 class_registry[name] = (full_module_name, name)
                 globals()[name] = obj
                 logger.info(f"Adding agent: {name}")
-                __all__.append(name) # type: ignore
+                __all__.append(name)  # type: ignore
     except ImportError as e:
         logger.error(f"Error importing {full_module_name}: {e}")
         raise e
diff --git a/src/utils/agents/base.py b/src/utils/agents/base.py
index 352da48..1b9173f 100644
--- a/src/utils/agents/base.py
+++ b/src/utils/agents/base.py
@@ -1,8 +1,17 @@
 from __future__ import annotations
-from pydantic import BaseModel, PrivateAttr, Field # type: ignore
+from pydantic import BaseModel, PrivateAttr, Field  # type: ignore
 from typing import (
-    Literal, get_args, List, AsyncGenerator, TYPE_CHECKING, Optional, ClassVar, Any,
-    TypeAlias, Dict, Tuple
+    Literal,
+    get_args,
+    List,
+    AsyncGenerator,
+    TYPE_CHECKING,
+    Optional,
+    ClassVar,
+    Any,
+    TypeAlias,
+    Dict,
+    Tuple,
 )
 import json
 import time
@@ -10,59 +19,66 @@ import inspect
 from abc import ABC
 import asyncio
 
-from prometheus_client import Counter, Summary, CollectorRegistry # type: ignore
+from prometheus_client import Counter, Summary, CollectorRegistry  # type: ignore
 
-from .. setup_logging import setup_logging
+from ..setup_logging import setup_logging
 
 logger = setup_logging()
 
 # Only import Context for type checking
 if TYPE_CHECKING:
-  from .. context import Context
+    from ..context import Context
 
-from . types import agent_registry
+from .types import agent_registry
 from .. import defines
-from .. message import Message, Tunables
-from .. metrics import Metrics
-from .. tools import ( TickerValue, WeatherForecast, AnalyzeSite, DateTime, llm_tools ) # type: ignore -- dynamically added to __all__
-from .. conversation import Conversation
+from ..message import Message, Tunables
+from ..metrics import Metrics
+from ..tools import TickerValue, WeatherForecast, AnalyzeSite, DateTime, llm_tools  # type: ignore -- dynamically added to __all__
+from ..conversation import Conversation
+
 
 class LLMMessage(BaseModel):
-    role : str = Field(default="")
-    content : str = Field(default="")
-    tool_calls : Optional[List[Dict]] = Field(default={}, exclude=True)
+    role: str = Field(default="")
+    content: str = Field(default="")
+    tool_calls: Optional[List[Dict]] = Field(default={}, exclude=True)
+
 
 class Agent(BaseModel, ABC):
     """
     Base class for all agent types.
     This class defines the common attributes and methods for all agent types.
     """
+
     # Agent management with pydantic
     agent_type: Literal["base"] = "base"
     _agent_type: ClassVar[str] = agent_type  # Add this for registration
 
-
     # Tunables (sets default for new Messages attached to this agent)
     tunables: Tunables = Field(default_factory=Tunables)
 
     # Agent properties
-    system_prompt: str # Mandatory
+    system_prompt: str  # Mandatory
     conversation: Conversation = Conversation()
     context_tokens: int = 0
-    context: Optional[Context] = Field(default=None, exclude=True) # Avoid circular reference, require as param, and prevent serialization
+    context: Optional[Context] = Field(
+        default=None, exclude=True
+    )  # Avoid circular reference, require as param, and prevent serialization
     metrics: Metrics = Field(default_factory=Metrics, exclude=True)
 
     # context_size is shared across all subclasses
-    _context_size: ClassVar[int] = int(defines.max_context * 0.5)    
+    _context_size: ClassVar[int] = int(defines.max_context * 0.5)
+
     @property
     def context_size(self) -> int:
         return Agent._context_size
-    
+
     @context_size.setter
     def context_size(self, value: int):
         Agent._context_size = value
 
-    def set_optimal_context_size(self, llm: Any, model: str, prompt: str, ctx_buffer=2048) -> int:
+    def set_optimal_context_size(
+        self, llm: Any, model: str, prompt: str, ctx_buffer=2048
+    ) -> int:
         # # Get more accurate token count estimate using tiktoken or similar
         # response = llm.generate(
         #     model=model,
@@ -73,31 +89,33 @@ class Agent(BaseModel, ABC):
         #     }  # Don't generate any tokens, just tokenize
         # )
         # # The prompt_eval_count gives you the token count of your input
-        # tokens = response.get("prompt_eval_count", 0)        
-        
+        # tokens = response.get("prompt_eval_count", 0)
+
         # Most models average 1.3-1.5 tokens per word
         word_count = len(prompt.split())
         tokens = int(word_count * 1.4)
-        
+
         # Add buffer for safety
         total_ctx = tokens + ctx_buffer
-    
+
         if total_ctx > self.context_size:
-            logger.info(f"Increasing context size from {self.context_size} to {total_ctx}")
+            logger.info(
+                f"Increasing context size from {self.context_size} to {total_ctx}"
+            )
 
         # Grow the context size if necessary
         self.context_size = max(self.context_size, total_ctx)
         # Use actual model maximum context size
         return self.context_size
-    
+
     # Class and pydantic model management
     def __init_subclass__(cls, **kwargs) -> None:
         """Auto-register subclasses"""
         super().__init_subclass__(**kwargs)
         # Register this class if it has an agent_type
-        if hasattr(cls, 'agent_type') and cls.agent_type != Agent._agent_type:
+        if hasattr(cls, "agent_type") and cls.agent_type != Agent._agent_type:
             agent_registry.register(cls.agent_type, cls)
-    
+
     # def __init__(self, *, context=context, **data):
     #     super().__init__(**data)
     #     self.set_context(context)
@@ -118,12 +136,12 @@ class Agent(BaseModel, ABC):
 
     def set_context(self, context: Context):
         object.__setattr__(self, "context", context)
-        
+
     # Agent methods
     def get_agent_type(self):
         return self._agent_type
-    
-    async def prepare_message(self, message:Message) -> AsyncGenerator[Message, None]:
+
+    async def prepare_message(self, message: Message) -> AsyncGenerator[Message, None]:
         """
         Prepare message with context information in message.preamble
         """
@@ -133,7 +151,7 @@ class Agent(BaseModel, ABC):
         with self.metrics.prepare_duration.labels(agent=self.agent_type).time():
             if not self.context:
                 raise ValueError("Context is not set for this agent.")
-            
+
             # Generate RAG content if enabled, based on the content
             rag_context = ""
             if message.tunables.enable_rag and message.prompt:
@@ -146,7 +164,7 @@ class Agent(BaseModel, ABC):
                         return
                     if message.status != "done":
                         yield message
-                    
+
                 if "rag" in message.metadata and message.metadata["rag"]:
                     for rag in message.metadata["rag"]:
                         for doc in rag["documents"]:
@@ -159,16 +177,23 @@ class Agent(BaseModel, ABC):
 
             if message.tunables.enable_context and self.context.user_resume:
                 message.preamble["resume"] = self.context.user_resume
-            
+
             message.system_prompt = self.system_prompt
             message.status = "done"
             yield message
 
         return
 
-    async def process_tool_calls(self, llm: Any, model: str, message: Message, tool_message: Any, messages: List[LLMMessage]) -> AsyncGenerator[Message, None]:
+    async def process_tool_calls(
+        self,
+        llm: Any,
+        model: str,
+        message: Message,
+        tool_message: Any,
+        messages: List[LLMMessage],
+    ) -> AsyncGenerator[Message, None]:
         logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
-        
+
         self.metrics.tool_count.labels(agent=self.agent_type).inc()
         with self.metrics.tool_duration.labels(agent=self.agent_type).time():
 
@@ -185,12 +210,14 @@ class Agent(BaseModel, ABC):
             for i, tool_call in enumerate(tool_message.tool_calls):
                 arguments = tool_call.function.arguments
                 tool = tool_call.function.name
-            
+
                 # Yield status update before processing each tool
-                message.response = f"Processing tool {i+1}/{len(tool_message.tool_calls)}: {tool}..."
+                message.response = (
+                    f"Processing tool {i+1}/{len(tool_message.tool_calls)}: {tool}..."
+                )
                 yield message
                 logger.info(f"LLM - {message.response}")
-                
+
                 # Process the tool based on its type
                 match tool:
                     case "TickerValue":
@@ -199,40 +226,48 @@ class Agent(BaseModel, ABC):
                             ret = None
                         else:
                             ret = TickerValue(ticker)
-                    
+
                     case "AnalyzeSite":
                         url = arguments.get("url")
-                        question = arguments.get("question", "what is the summary of this content?")
-                    
+                        question = arguments.get(
+                            "question", "what is the summary of this content?"
+                        )
+
                         # Additional status update for long-running operations
-                        message.response = f"Retrieving and summarizing content from {url}..."
+                        message.response = (
+                            f"Retrieving and summarizing content from {url}..."
+                        )
                         yield message
-                        ret = await AnalyzeSite(llm=llm, model=model, url=url, question=question)
-                        
+                        ret = await AnalyzeSite(
+                            llm=llm, model=model, url=url, question=question
+                        )
+
                     case "DateTime":
                         tz = arguments.get("timezone")
                         ret = DateTime(tz)
-                    
+
                     case "WeatherForecast":
                         city = arguments.get("city")
                         state = arguments.get("state")
 
-                        message.response = f"Fetching weather data for {city}, {state}..."
-                        yield message                  
+                        message.response = (
+                            f"Fetching weather data for {city}, {state}..."
+                        )
+                        yield message
                         ret = WeatherForecast(city, state)
-                        
+
                     case _:
                         ret = None
-                
+
                 # Build response for this tool
                 tool_response = {
                     "role": "tool",
                     "content": json.dumps(ret),
-                    "name": tool_call.function.name
+                    "name": tool_call.function.name,
                 }
 
                 tool_metadata["tool_calls"].append(tool_response)
-                
+
                 if len(tool_metadata["tool_calls"]) == 0:
                     message.status = "done"
                     yield message
@@ -241,13 +276,15 @@ class Agent(BaseModel, ABC):
                 message_dict = LLMMessage(
                     role=tool_message.get("role", "assistant"),
                     content=tool_message.get("content", ""),
-                    tool_calls=[ {
-                        "function": {
-                            "name": tc["function"]["name"], 
-                            "arguments": tc["function"]["arguments"]
+                    tool_calls=[
+                        {
+                            "function": {
+                                "name": tc["function"]["name"],
+                                "arguments": tc["function"]["arguments"],
+                            }
                         }
-                        } for tc in tool_message.tool_calls
-                    ]
+                        for tc in tool_message.tool_calls
+                    ],
                 )
 
                 messages.append(message_dict)
@@ -261,8 +298,8 @@ class Agent(BaseModel, ABC):
             message.response = ""
             start_time = time.perf_counter()
             for response in llm.chat(
-                model=model, 
-                messages=messages, 
+                model=model,
+                messages=messages,
                 options={
                     **message.metadata["options"],
                     # "temperature": 0.5,
@@ -282,45 +319,62 @@ class Agent(BaseModel, ABC):
                     message.metadata["eval_count"] += response.eval_count
                     message.metadata["eval_duration"] += response.eval_duration
                     message.metadata["prompt_eval_count"] += response.prompt_eval_count
-                    message.metadata["prompt_eval_duration"] += response.prompt_eval_duration
-                    self.context_tokens = response.prompt_eval_count + response.eval_count
+                    message.metadata[
+                        "prompt_eval_duration"
+                    ] += response.prompt_eval_duration
+                    self.context_tokens = (
+                        response.prompt_eval_count + response.eval_count
+                    )
                     message.status = "done"
                     yield message
 
             end_time = time.perf_counter()
-            message.metadata["timers"]["llm_with_tools"] = f"{(end_time - start_time):.4f}"
+            message.metadata["timers"][
+                "llm_with_tools"
+            ] = f"{(end_time - start_time):.4f}"
             return
-            
+
     def collect_metrics(self, response):
-        self.metrics.tokens_prompt.labels(agent=self.agent_type).inc(response.prompt_eval_count)
+        self.metrics.tokens_prompt.labels(agent=self.agent_type).inc(
+            response.prompt_eval_count
+        )
         self.metrics.tokens_eval.labels(agent=self.agent_type).inc(response.eval_count)
 
-    async def generate_llm_response(self, llm: Any, model: str, message: Message, temperature = 0.7) -> AsyncGenerator[Message, None]:
+    async def generate_llm_response(
+        self, llm: Any, model: str, message: Message, temperature=0.7
+    ) -> AsyncGenerator[Message, None]:
         logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
-        
+
         self.metrics.generate_count.labels(agent=self.agent_type).inc()
         with self.metrics.generate_duration.labels(agent=self.agent_type).time():
             if not self.context:
                 raise ValueError("Context is not set for this agent.")
 
-            # Create a pruned down message list based purely on the prompt and responses, 
+            # Create a pruned down message list based purely on the prompt and responses,
             # discarding the full preamble generated by prepare_message
-            messages: List[LLMMessage] = [ LLMMessage(role="system", content=message.system_prompt) ]
-            messages.extend([
-                item for m in self.conversation
-                for item in [
-                    LLMMessage(role="user", content=m.prompt.strip()),
-                    LLMMessage(role="assistant", content=m.response.strip())
+            messages: List[LLMMessage] = [
+                LLMMessage(role="system", content=message.system_prompt)
+            ]
+            messages.extend(
+                [
+                    item
+                    for m in self.conversation
+                    for item in [
+                        LLMMessage(role="user", content=m.prompt.strip()),
+                        LLMMessage(role="assistant", content=m.response.strip()),
+                    ]
                 ]
-            ])
+            )
             # Only the actual user query is provided with the full context message
-            messages.append(LLMMessage(role="user", content=message.context_prompt.strip()))
+            messages.append(
+                LLMMessage(role="user", content=message.context_prompt.strip())
+            )
 
-            #message.metadata["messages"] = messages
-            message.metadata["options"]={
+            # message.metadata["messages"] = messages
+            message.metadata["options"] = {
                 "seed": 8911,
                 "num_ctx": self.context_size,
-                "temperature": temperature # Higher temperature to encourage tool usage
+                "temperature": temperature,  # Higher temperature to encourage tool usage
             }
 
             # Create a dict for storing various timing stats
@@ -329,7 +383,7 @@ class Agent(BaseModel, ABC):
             use_tools = message.tunables.enable_tools and len(self.context.tools) > 0
             message.metadata["tools"] = {
                 "available": llm_tools(self.context.tools),
-                "used": False
+                "used": False,
             }
             tool_metadata = message.metadata["tools"]
 
@@ -355,14 +409,16 @@ class Agent(BaseModel, ABC):
                     tools=tool_metadata["available"],
                     options={
                         **message.metadata["options"],
-                        #"num_predict": 1024, # "Low" token limit to cut off after tool call
+                        # "num_predict": 1024, # "Low" token limit to cut off after tool call
                     },
-                    stream=False  # No need to stream the probe
+                    stream=False,  # No need to stream the probe
                 )
                 self.collect_metrics(response)
 
                 end_time = time.perf_counter()
-                message.metadata["timers"]["tool_check"] = f"{(end_time - start_time):.4f}"
+                message.metadata["timers"][
+                    "tool_check"
+                ] = f"{(end_time - start_time):.4f}"
                 if not response.message.tool_calls:
                     logger.info("LLM indicates tools will not be used")
                     # The LLM will not use tools, so disable use_tools so we can stream the full response
@@ -374,7 +430,9 @@ class Agent(BaseModel, ABC):
                 logger.info("LLM indicates tools will be used")
 
                 # Tools are enabled and available and the LLM indicated it will use them
-                message.response = f"Performing tool analysis step 2/2 (tool use suspected)..."
+                message.response = (
+                    f"Performing tool analysis step 2/2 (tool use suspected)..."
+                )
                 yield message
 
                 logger.info(f"Performing LLM call with tools")
@@ -384,14 +442,16 @@ class Agent(BaseModel, ABC):
                     messages=tool_metadata["messages"],  # messages,
                     tools=tool_metadata["available"],
                     options={
-                    **message.metadata["options"],
+                        **message.metadata["options"],
                     },
-                    stream=False
+                    stream=False,
                 )
                 self.collect_metrics(response)
 
                 end_time = time.perf_counter()
-                message.metadata["timers"]["non_streaming"] = f"{(end_time - start_time):.4f}"
+                message.metadata["timers"][
+                    "non_streaming"
+                ] = f"{(end_time - start_time):.4f}"
 
                 if not response:
                     message.status = "error"
@@ -403,16 +463,24 @@ class Agent(BaseModel, ABC):
                     tool_metadata["used"] = response.message.tool_calls
                     # Process all yielded items from the handler
                     start_time = time.perf_counter()
-                    async for message in self.process_tool_calls(llm=llm, model=model, message=message, tool_message=response.message, messages=messages):
+                    async for message in self.process_tool_calls(
+                        llm=llm,
+                        model=model,
+                        message=message,
+                        tool_message=response.message,
+                        messages=messages,
+                    ):
                         if message.status == "error":
                             yield message
                             return
                         yield message
                     end_time = time.perf_counter()
-                    message.metadata["timers"]["process_tool_calls"] = f"{(end_time - start_time):.4f}"
+                    message.metadata["timers"][
+                        "process_tool_calls"
+                    ] = f"{(end_time - start_time):.4f}"
                     message.status = "done"
                     return
-                
+
                 logger.info("LLM indicated tools will be used, and then they weren't")
                 message.response = response.message.content
                 message.status = "done"
@@ -427,8 +495,8 @@ class Agent(BaseModel, ABC):
             message.response = ""
             start_time = time.perf_counter()
             for response in llm.chat(
-                model=model, 
-                messages=messages, 
+                model=model,
+                messages=messages,
                 options={
                     **message.metadata["options"],
                 },
@@ -452,16 +520,22 @@ class Agent(BaseModel, ABC):
                     message.metadata["eval_count"] += response.eval_count
                     message.metadata["eval_duration"] += response.eval_duration
                     message.metadata["prompt_eval_count"] += response.prompt_eval_count
-                    message.metadata["prompt_eval_duration"] += response.prompt_eval_duration
-                    self.context_tokens = response.prompt_eval_count + response.eval_count
+                    message.metadata[
+                        "prompt_eval_duration"
+                    ] += response.prompt_eval_duration
+                    self.context_tokens = (
+                        response.prompt_eval_count + response.eval_count
+                    )
                     message.status = "done"
                     yield message
 
             end_time = time.perf_counter()
             message.metadata["timers"]["streamed"] = f"{(end_time - start_time):.4f}"
         return
-    
-    async def process_message(self, llm: Any, model: str, message:Message) -> AsyncGenerator[Message, None]:
+
+    async def process_message(
+        self, llm: Any, model: str, message: Message
+    ) -> AsyncGenerator[Message, None]:
         logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
 
         self.metrics.process_count.labels(agent=self.agent_type).inc()
@@ -469,23 +543,31 @@ class Agent(BaseModel, ABC):
 
             if not self.context:
                 raise ValueError("Context is not set for this agent.")
-            
-            logger.info("TODO: Implement delay queing; busy for same agent, otherwise return queue size and estimated wait time")
-            spinner: List[str] = ['\\', '|', '/', '-']
-            tick : int = 0
+
+            logger.info(
+                "TODO: Implement delay queing; busy for same agent, otherwise return queue size and estimated wait time"
+            )
+            spinner: List[str] = ["\\", "|", "/", "-"]
+            tick: int = 0
             while self.context.processing:
                 message.status = "waiting"
-                message.response = f"Busy processing another request. Please wait. {spinner[tick]}"
+                message.response = (
+                    f"Busy processing another request. Please wait. {spinner[tick]}"
+                )
                 tick = (tick + 1) % len(spinner)
                 yield message
                 await asyncio.sleep(1)  # Allow the event loop to process the write
 
             self.context.processing = True
 
-            message.metadata["system_prompt"] = f"<|system|>\n{self.system_prompt.strip()}\n</|system|>"
+            message.metadata["system_prompt"] = (
+                f"<|system|>\n{self.system_prompt.strip()}\n</|system|>"
+            )
             message.context_prompt = ""
             for p in message.preamble.keys():
-                message.context_prompt += f"\n<|{p}|>\n{message.preamble[p].strip()}\n</|{p}>\n\n"
+                message.context_prompt += (
+                    f"\n<|{p}|>\n{message.preamble[p].strip()}\n</|{p}>\n\n"
+                )
             message.context_prompt += f"{message.prompt}"
 
             # Estimate token length of new messages
@@ -493,20 +575,24 @@ class Agent(BaseModel, ABC):
             message.status = "thinking"
             yield message
 
-            message.metadata["context_size"] = self.set_optimal_context_size(llm, model, prompt=message.context_prompt)
-            
+            message.metadata["context_size"] = self.set_optimal_context_size(
+                llm, model, prompt=message.context_prompt
+            )
+
             message.response = f"Processing {'RAG augmented ' if message.metadata['rag'] else ''}query..."
             message.status = "thinking"
             yield message
 
-            async for message in self.generate_llm_response(llm=llm, model=model, message=message):
+            async for message in self.generate_llm_response(
+                llm=llm, model=model, message=message
+            ):
                 # logger.info(f"LLM: {message.status} - {f'...{message.response[-20:]}' if len(message.response) > 20 else message.response}")
                 if message.status == "error":
                     yield message
                     self.context.processing = False
                     return
                 yield message
-            
+
             # Done processing, add message to conversation
             message.status = "done"
             self.conversation.add(message)
@@ -514,6 +600,6 @@ class Agent(BaseModel, ABC):
 
         return
 
+
 # Register the base agent
 agent_registry.register(Agent._agent_type, Agent)
-
diff --git a/src/utils/agents/chat.py b/src/utils/agents/chat.py
index 6a6fdfc..84c7b09 100644
--- a/src/utils/agents/chat.py
+++ b/src/utils/agents/chat.py
@@ -3,9 +3,10 @@ from typing import Literal, AsyncGenerator, ClassVar, Optional, Any
 from datetime import datetime
 import inspect
 
-from . base import Agent, agent_registry
-from .. message import Message
-from .. setup_logging import setup_logging
+from .base import Agent, agent_registry
+from ..message import Message
+from ..setup_logging import setup_logging
+
 logger = setup_logging()
 
 system_message = f"""
@@ -26,35 +27,42 @@ When answering queries, follow these steps:
 Always use tools, <|resume|>, and <|context|> when possible. Be concise, and never make up information. If you do not know the answer, say so.
 """
 
+
 class Chat(Agent):
-  """
-  Chat Agent
-  """
-  agent_type: Literal["chat"] = "chat" # type: ignore
-  _agent_type: ClassVar[str] = agent_type  # Add this for registration
+    """
+    Chat Agent
+    """
 
-  system_prompt: str = system_message
+    agent_type: Literal["chat"] = "chat"  # type: ignore
+    _agent_type: ClassVar[str] = agent_type  # Add this for registration
 
-  async def prepare_message(self, message:Message) -> AsyncGenerator[Message, None]:
-    logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
-    if not self.context:
-      raise ValueError("Context is not set for this agent.")
-    
-    async for message in super().prepare_message(message):
-      if message.status != "done":
-        yield message
+    system_prompt: str = system_message
 
-    if message.preamble:
-      excluded = {}
-      preamble_types = [f"<|{p}|>" for p in message.preamble.keys() if p not in excluded]
-      preamble_types_AND = " and ".join(preamble_types)
-      preamble_types_OR = " or ".join(preamble_types)
-      message.preamble["rules"] = f"""\
+    async def prepare_message(self, message: Message) -> AsyncGenerator[Message, None]:
+        logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
+        if not self.context:
+            raise ValueError("Context is not set for this agent.")
+
+        async for message in super().prepare_message(message):
+            if message.status != "done":
+                yield message
+
+        if message.preamble:
+            excluded = {}
+            preamble_types = [
+                f"<|{p}|>" for p in message.preamble.keys() if p not in excluded
+            ]
+            preamble_types_AND = " and ".join(preamble_types)
+            preamble_types_OR = " or ".join(preamble_types)
+            message.preamble[
+                "rules"
+            ] = f"""\
 - Answer the question based on the information provided in the {preamble_types_AND} sections by incorporate it seamlessly and refer to it using natural language instead of mentioning {preamble_types_OR} or quoting it directly.
 - If there is no information in these sections, answer based on your knowledge, or use any available tools.
 - Avoid phrases like 'According to the {preamble_types[0]}' or similar references to the {preamble_types_OR}.
 """
-    message.preamble["question"] = "Respond to:"
+        message.preamble["question"] = "Respond to:"
+
 
 # Register the base agent
 agent_registry.register(Chat._agent_type, Chat)
diff --git a/src/utils/agents/fact_check.py b/src/utils/agents/fact_check.py
index eb5d754..5131aa9 100644
--- a/src/utils/agents/fact_check.py
+++ b/src/utils/agents/fact_check.py
@@ -1,13 +1,21 @@
 from __future__ import annotations
-from pydantic import model_validator # type: ignore
-from typing import Literal, ClassVar, Optional, Any, AsyncGenerator, List # NOTE: You must import Optional for late binding to work
+from pydantic import model_validator  # type: ignore
+from typing import (
+    Literal,
+    ClassVar,
+    Optional,
+    Any,
+    AsyncGenerator,
+    List,
+)  # NOTE: You must import Optional for late binding to work
 from datetime import datetime
 import inspect
 
-from . base import Agent, agent_registry
-from .. conversation import Conversation
-from .. message import Message
-from .. setup_logging import setup_logging
+from .base import Agent, agent_registry
+from ..conversation import Conversation
+from ..message import Message
+from ..setup_logging import setup_logging
+
 logger = setup_logging()
 
 system_fact_check = f"""
@@ -21,50 +29,56 @@ When answering queries, follow these steps:
 - Avoid phrases like 'According to the <|context|>' or similar references to the <|context|>, <|generated-resume|>, or <|resume|> tags.
 """.strip()
 
+
 class FactCheck(Agent):
-  agent_type: Literal["fact_check"] = "fact_check" # type: ignore
-  _agent_type: ClassVar[str] = agent_type  # Add this for registration
+    agent_type: Literal["fact_check"] = "fact_check"  # type: ignore
+    _agent_type: ClassVar[str] = agent_type  # Add this for registration
 
-  system_prompt:str = system_fact_check
-  facts: str
+    system_prompt: str = system_fact_check
+    facts: str
 
-  @model_validator(mode="after")
-  def validate_facts(self):
-    if not self.facts.strip():
-      raise ValueError("Facts cannot be empty")
-    return self
+    @model_validator(mode="after")
+    def validate_facts(self):
+        if not self.facts.strip():
+            raise ValueError("Facts cannot be empty")
+        return self
 
-  async def prepare_message(self, message:Message) -> AsyncGenerator[Message, None]:
-    logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
-    if not self.context:
-      raise ValueError("Context is not set for this agent.")
+    async def prepare_message(self, message: Message) -> AsyncGenerator[Message, None]:
+        logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
+        if not self.context:
+            raise ValueError("Context is not set for this agent.")
 
-    resume_agent = self.context.get_agent("resume")
-    if not resume_agent:
-      raise ValueError("resume agent does not exist")
+        resume_agent = self.context.get_agent("resume")
+        if not resume_agent:
+            raise ValueError("resume agent does not exist")
 
-    message.tunables.enable_tools = False
+        message.tunables.enable_tools = False
 
-    async for message in super().prepare_message(message):
-      if message.status != "done":
-        yield message
+        async for message in super().prepare_message(message):
+            if message.status != "done":
+                yield message
 
-    message.preamble["generated-resume"] = resume_agent.resume
-    message.preamble["discrepancies"] = self.facts
-    
-    excluded = {"job_description"}
-    preamble_types = [f"<|{p}|>" for p in message.preamble.keys() if p not in excluded]
-    preamble_types_AND = " and ".join(preamble_types)
-    preamble_types_OR = " or ".join(preamble_types)
-    message.preamble["rules"] = f"""\
+        message.preamble["generated-resume"] = resume_agent.resume
+        message.preamble["discrepancies"] = self.facts
+
+        excluded = {"job_description"}
+        preamble_types = [
+            f"<|{p}|>" for p in message.preamble.keys() if p not in excluded
+        ]
+        preamble_types_AND = " and ".join(preamble_types)
+        preamble_types_OR = " or ".join(preamble_types)
+        message.preamble[
+            "rules"
+        ] = f"""\
 - Answer the question based on the information provided in the {preamble_types_AND} sections by incorporate it seamlessly and refer to it using natural language instead of mentioning {preamble_types_OR} or quoting it directly.
 - If there is no information in these sections, answer based on your knowledge, or use any available tools.
 - Avoid phrases like 'According to the {preamble_types[0]}' or similar references to the {preamble_types_OR}.
 """
-    message.preamble["question"] = "Respond to:"
+        message.preamble["question"] = "Respond to:"
+
+        yield message
+        return
 
-    yield message
-    return
 
 # Register the base agent
 agent_registry.register(FactCheck._agent_type, FactCheck)
diff --git a/src/utils/agents/job_description.py b/src/utils/agents/job_description.py
index cd9f428..b2e09ba 100644
--- a/src/utils/agents/job_description.py
+++ b/src/utils/agents/job_description.py
@@ -1,6 +1,14 @@
 from __future__ import annotations
-from pydantic import model_validator, Field # type: ignore
-from typing import Dict, Literal, ClassVar, Optional, Any, AsyncGenerator, List # NOTE: You must import Optional for late binding to work
+from pydantic import model_validator, Field  # type: ignore
+from typing import (
+    Dict,
+    Literal,
+    ClassVar,
+    Optional,
+    Any,
+    AsyncGenerator,
+    List,
+)  # NOTE: You must import Optional for late binding to work
 from datetime import datetime
 import inspect
 import re
@@ -9,345 +17,768 @@ import traceback
 import asyncio
 import time
 from collections import defaultdict
+import asyncio
+import numpy as np  # type: ignore
+
+from .base import Agent, agent_registry, LLMMessage
+from ..conversation import Conversation
+from ..message import Message
+from ..setup_logging import setup_logging
 
-from . base import Agent, agent_registry, LLMMessage
-from .. conversation import Conversation
-from .. message import Message
-from .. setup_logging import setup_logging
 logger = setup_logging()
 
 
 system_generate_resume = """
-You are an objective skills analyzer for resume tailoring. Your task is to 
-analyze a job description and a candidate's background, identifying ONLY 
-legitimate matches without fabrication. Your analysis will be used in a 
-second stage to generate a tailored resume.
-
-## INSTRUCTIONS:
-
-1. Analyze the job description to extract:
-   - Required skills and technologies
-   - Required experience types
-   - Company values and culture indicators
-   - Primary responsibilities of the role
-
-2. Analyze the candidate's resume and context to identify:
-   - Actual skills and technologies the candidate possesses
-   - Actual experience types the candidate has
-   - Relevant achievements that align with the job
-
-3. Create a structured output with these sections:
-
-### OUTPUT FORMAT:
-
-```json
-{
-  "job_requirements": {
-    "required_skills": ["skill1", "skill2"...],
-    "preferred_skills": ["skill1", "skill2"...],
-    "required_experience": ["exp1", "exp2"...],
-    "company_values": ["value1", "value2"...]
-  },
-  "candidate_qualifications": {
-    "matching_skills": ["skill1", "skill2"...],
-    "transferable_skills": ["skill1", "skill2"...],
-    "missing_skills": ["skill1", "skill2"...],
-    "relevant_experience": [
-      {
-        "role": "Role title",
-        "relevance": "Explanation of how this role is relevant"
-      }
-    ],
-    "relevant_achievements": [
-      {
-        "achievement": "Achievement description",
-        "relevance": "Explanation of relevance to job"
-      }
-    ]
-  },
-  "emphasis_recommendations": {
-    "highlight_skills": ["skill1", "skill2"...],
-    "highlight_experiences": ["exp1", "exp2"...],
-    "transferable_narratives": [
-      "Description of how candidate's experience in X transfers to required skill Y"
-    ],
-    "honest_gap_acknowledgment": [
-      "Skills/technologies the candidate lacks with no reasonable transfer claim"
-    ]
-  }
-}
-```
-
-### CRITICAL RULES:
-
-1. DO NOT INVENT OR IMPLY ANY SKILLS OR EXPERIENCE not explicitly present in the candidate's materials
-2. A skill is only "matching" if the EXACT SAME technology/skill appears in both job requirements and candidate background
-3. "Transferable skills" must have a legitimate connection - don't stretch credibility
-4. Be brutally honest in "missing_skills" and "honest_gap_acknowledgment" sections
-5. For each skill or experience match, note the specific section/line from the candidate's materials where it appears
-6. If a job requirement has no legitimate match or transfer in the candidate's background, it MUST appear in "missing_skills"
-
-The output of this analysis will be used to create a truthful, tailored resume. Accuracy is critical.
 """.strip()
 
 system_job_description = f"""
-Launched on {datetime.now().isoformat()}.
-
-You are a hiring and job placing specialist. Your task is to answers about a job description.
-
-When answering queries, follow these steps:
-- Analyze the <|job_description|> to provide insights for the asked question.
-- If any financial information is requested, be sure to account for inflation.
-- If there is information in the <|context|>, <|job_description|>, <|context|>, or <|resume|> sections to enhance the answer, incorporate it seamlessly and refer to it using natural language instead of mentioning '<|job_description|>' (etc.) or quoting it directly.
-- Avoid phrases like 'According to the <|context|>' or similar references to the <|context|>, <|job_description|>, <|resume|>, or <|context|> tags.
 """.strip()
 
 system_user_qualifications = """
-You are an objective job requirements analyzer. Your task is to extract and categorize the specific skills, experiences, and qualifications required in a job description WITHOUT any reference to any candidate.
-
-## INSTRUCTIONS:
-
-1. Analyze ONLY the job description provided.
-2. Extract and categorize all requirements and preferences.
-3. DO NOT consider any candidate information - this is a pure job analysis task.
-
-## OUTPUT FORMAT:
-
-```json
-{
-  "job_requirements": {
-    "technical_skills": {
-      "required": ["skill1", "skill2"...],
-      "preferred": ["skill1", "skill2"...]
-    },
-    "experience_requirements": {
-      "required": ["exp1", "exp2"...],
-      "preferred": ["exp1", "exp2"...]
-    },
-    "education_requirements": ["req1", "req2"...],
-    "soft_skills": ["skill1", "skill2"...],
-    "industry_knowledge": ["knowledge1", "knowledge2"...],
-    "responsibilities": ["resp1", "resp2"...],
-    "company_values": ["value1", "value2"...]
-  }
-}
-```
-
-Be specific and detailed in your extraction. Break down compound requirements into individual components.
-For example, "5+ years experience with React, Node.js and MongoDB" should be separated into:
-- Experience: "5+ years software development"
-- Technical skills: "React", "Node.js", "MongoDB"
-
-Avoid vague categorizations and be precise about whether skills are explicitly required or just preferred.
-</|system|>
-
-<|job_description|>
-[INSERT JOB DESCRIPTION HERE]
-</|job_description|>
+Answer questions about the job description.
 """
 
+
 class JobDescription(Agent):
-  agent_type: Literal["job_description"] = "job_description" # type: ignore
-  _agent_type: ClassVar[str] = agent_type  # Add this for registration
+    agent_type: Literal["job_description"] = "job_description"  # type: ignore
+    _agent_type: ClassVar[str] = agent_type  # Add this for registration
 
-  system_prompt: str = system_generate_resume
-  job_description: str
+    system_prompt: str = system_generate_resume
+    job_description: str
 
-  llm: Any = Field(default=None, exclude=True)
-  model: str = Field(default=None, exclude=True)
+    llm: Any = Field(default=None, exclude=True)
+    model: str = Field(default=None, exclude=True)
 
-  @model_validator(mode="after")
-  def validate_job_description(self):
-    if not self.job_description.strip():
-      raise ValueError("Job description cannot be empty")
-    return self
+    @model_validator(mode="after")
+    def validate_job_description(self):
+        if not self.job_description.strip():
+            raise ValueError("Job description cannot be empty")
+        return self
 
-  async def prepare_message(self, message:Message) -> AsyncGenerator[Message, None]:
-    logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
-    if not self.context:
-      raise ValueError("Context is not set for this agent.")
-    
-    # async for message in super().prepare_message(message):
-    #   if message.status != "done":
-    #     yield message
-    # Always add the job description, user resume, and question
-    message.preamble["job_description"] = self.job_description
-    message.preamble["resume"] = self.context.user_resume
-    yield message
-    return
+    async def prepare_message(self, message: Message) -> AsyncGenerator[Message, None]:
+        logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
+        if not self.context:
+            raise ValueError("Context is not set for this agent.")
 
-  async def process_message(self, llm: Any, model: str, message:Message) -> AsyncGenerator[Message, None]:
-    logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
-    if not self.context:
-      raise ValueError("Context is not set for this agent.")
-    
-    # Generating a resume should not use any tools
-    message.tunables.enable_tools = False
-
-    original_prompt = message.prompt
-
-    logger.info("TODO: Implement delay queing; busy for same agent, otherwise return queue size and estimated wait time")
-    spinner: List[str] = ['\\', '|', '/', '-']
-    tick : int = 0
-    while self.context.processing:
-        message.status = "waiting"
-        message.response = f"Busy processing another request. Please wait. {spinner[tick]}"
-        tick = (tick + 1) % len(spinner)
-        yield message
-        await asyncio.sleep(1)  # Allow the event loop to process the write
-
-    self.context.processing = True
-
-    original_message = message.model_copy()
-
-    self.llm = llm
-    self.model = model
-    self.metrics.generate_count.labels(agent=self.agent_type).inc()
-    with self.metrics.generate_duration.labels(agent=self.agent_type).time():
-      job_description = message.preamble["job_description"]
-      resume = message.preamble["resume"]
-      
-      try:
-        async for message in self.generate_factual_tailored_resume(message=message, job_description=job_description, resume=resume):
-          if message.status != "done":
-            yield message
-        message.prompt = original_message.prompt
-        yield message
-
-      except Exception as e:
-        message.status = "error"
-        logger.error(message.response)
-        message.response = f"Error in resume generation process: {str(e)}"
-        logger.error(message.response)
-        logger.error(traceback.format_exc())
+        # async for message in super().prepare_message(message):
+        #   if message.status != "done":
+        #     yield message
+        # Always add the job description, user resume, and question
+        message.preamble["job_description"] = self.job_description
+        message.preamble["resume"] = self.context.user_resume
         yield message
         return
-    
-    # Done processing, add message to conversation
-    message.status = "done"
-    self.conversation.add(message)
-    self.context.processing = False
 
-    # Add the "Job requirements" message
-    if "generate_factual_tailored_resume" in message.metadata and "job_requirements" in message.metadata["generate_factual_tailored_resume"]:
-      new_message = Message(prompt="Show job requirements")
-      job_requirements = message.metadata["generate_factual_tailored_resume"]["job_requirements"]["results"]
-      new_message.response = f"```json\n\n{json.dumps(job_requirements, indent=2)}\n```\n"
-      new_message.status = "done"
-      new_message.actions = [ "job_description" ]
-      self.conversation.add(new_message)
+    async def process_message(
+        self, llm: Any, model: str, message: Message
+    ) -> AsyncGenerator[Message, None]:
+        logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
+        if not self.context:
+            raise ValueError("Context is not set for this agent.")
 
-    self.system_prompt = system_user_qualifications
+        # Generating a resume should not use any tools
+        message.tunables.enable_tools = False
 
-    resume_agent = self.context.get_agent(agent_type="resume")
-    fact_check_agent = self.context.get_agent(agent_type="fact_check")
-    if not resume_agent:
-      # Add the "Generated Resume" message
-      if "generate_factual_tailored_resume" in message.metadata and "analyze_candidate_qualifications" in message.metadata["generate_factual_tailored_resume"]:
-        resume_agent = self.context.get_or_create_agent(agent_type="resume", resume=message.response)
-        resume_message = Message(prompt="Show candidate qualifications")
-        qualifications = message.metadata["generate_factual_tailored_resume"]["analyze_candidate_qualifications"]["results"]
-        resume_message.response = f"```json\n\n{json.dumps(qualifications, indent=2)}\n```\n"
-        resume_message.status = "done"
-        resume_agent.conversation.add(resume_message)
+        logger.info(
+            "TODO: Implement delay queing; busy for same agent, otherwise return queue size and estimated wait time"
+        )
+        spinner: List[str] = ["\\", "|", "/", "-"]
+        tick: int = 0
+        while self.context.processing:
+            message.status = "waiting"
+            message.response = (
+                f"Busy processing another request. Please wait. {spinner[tick]}"
+            )
+            tick = (tick + 1) % len(spinner)
+            yield message
+            await asyncio.sleep(1)  # Allow the event loop to process the write
 
+        self.context.processing = True
+
+        original_message = message.model_copy()
+        original_message.response = ""
+        self.conversation.add(original_message)
+
+        self.llm = llm
+        self.model = model
+        self.metrics.generate_count.labels(agent=self.agent_type).inc()
+        with self.metrics.generate_duration.labels(agent=self.agent_type).time():
+            job_description = message.preamble["job_description"]
+            resume = message.preamble["resume"]
+
+            try:
+                async for message in self.generate_factual_tailored_resume(
+                    message=message, job_description=job_description, resume=resume
+                ):
+                    if message.status != "done":
+                        yield message
+                message.status = "done"
+                yield message
+
+            except Exception as e:
+                message.status = "error"
+                logger.error(message.response)
+                message.response = f"Error in resume generation process: {str(e)}"
+                logger.error(message.response)
+                logger.error(traceback.format_exc())
+                yield message
+                return
+
+        # Done processing, add message to conversation
+        self.context.processing = False
+
+        resume_generation = message.metadata.get("resume_generation", {})
+        if not resume_generation:
+            message.response = (
+                "Generation did not generate metadata necessary for processing."
+            )
+            message.status = "error"
+            yield message
+            return
+
+        generated_resume = resume_generation.get(
+            "resume", "Resume generation resulted in an empty resume."
+        )
+
+        facts = resume_generation.get("match_stats", {"results": "No skills matched"})
+
+        # Add the "Job requirements" message
+
+        self.system_prompt = system_user_qualifications
+
+        resume_agent = self.context.get_agent(agent_type="resume")
+        if not resume_agent:
+            resume_agent = self.context.get_or_create_agent(
+                agent_type="resume", resume=generated_resume
+            )
+        else:
+            resume_agent.conversation.reset()
         resume_message = message.model_copy()
-        resume_message.prompt = "Show generated resume"
-        resume_message.response = message.response
+        resume_message.prompt = ""
+        resume_message.response = generated_resume
         resume_message.status = "done"
         resume_agent.conversation.add(resume_message)
 
-        message.response = "Resume generated."
         message.actions.append("resume_generated")
 
-      # Add the "Fact Check" message
-      if "generate_factual_tailored_resume" in message.metadata and "verify_resume" in message.metadata["generate_factual_tailored_resume"]:
-        if "second_pass" in message.metadata["generate_factual_tailored_resume"]["verify_resume"]:
-          verification = message.metadata["generate_factual_tailored_resume"]["verify_resume"]["second_pass"]["results"]
+        fact_check_agent = self.context.get_agent(agent_type="fact_check")
+        if not fact_check_agent:
+            fact_check_agent = self.context.get_or_create_agent(
+                agent_type="fact_check", facts=json.dumps(facts, indent=2)
+            )
         else:
-          verification = message.metadata["generate_factual_tailored_resume"]["verify_resume"]["first_pass"]["results"]
+            fact_check_agent.conversation.reset()
 
-        fact_check_agent = self.context.get_or_create_agent(agent_type="fact_check", facts=json.dumps(verification, indent=2))
         fact_check_message = message.model_copy()
-        fact_check_message.prompt = "Show verification"
-        fact_check_message.response = f"```json\n\n{json.dumps(verification, indent=2)}\n```\n"
+        fact_check_message.prompt = ""
+        clean_json = re.sub(r"```", "'''", json.dumps(facts, indent=2))
+        fact_check_message.response = f"```json\n\n{clean_json}\n```\n"
         fact_check_message.status = "done"
         fact_check_agent.conversation.add(fact_check_message)
 
-        message.prompt = original_prompt
-        message.response = "Resume generated and verified."
         message.actions.append("facts_checked")
-  
-    # Return the final message
-    yield message
-    return
 
-  # BEGIN
+        # Return the final message
+        message.response = "Resume generated."
+        yield message
+        return
 
-  # Helper functions
-  def extract_json_from_text(self, text: str) -> str:
-      """Extract JSON string from text that may contain other content."""
-      json_pattern = r'```json\s*([\s\S]*?)\s*```'
-      match = re.search(json_pattern, text)
-      if match:
-          return match.group(1).strip()
-      
-      # Try to find JSON without the markdown code block
-      json_pattern = r'({[\s\S]*})'
-      match = re.search(json_pattern, text)
-      if match:
-          return match.group(1).strip()
-      
-      raise ValueError("No JSON found in the response")
+    # BEGIN
 
-  def validate_job_requirements(self, job_requirements: Dict) -> None:
-      """Validate the structure of job requirements."""
-      required_keys = ["job_requirements"]
-      
-      if not all(key in job_requirements for key in required_keys):
-          missing = [key for key in required_keys if key not in job_requirements]
-          raise ValueError(f"Missing required keys in job requirements: {missing}")
-      
-      # Additional validation can be added here
+    # Helper functions
+    def extract_json_from_text(self, text: str) -> str:
+        """Extract JSON string from text that may contain other content."""
+        json_pattern = r"```json\s*([\s\S]*?)\s*```"
+        match = re.search(json_pattern, text)
+        if match:
+            return match.group(1).strip()
 
-  def validate_candidate_qualifications(self, candidate_qualifications: Dict) -> None:
-      """Validate the structure of candidate qualifications."""
-      required_keys = ["candidate_qualifications"]
-      
-      if not all(key in candidate_qualifications for key in required_keys):
-          missing = [key for key in required_keys if key not in candidate_qualifications]
-          raise ValueError(f"Missing required keys in candidate qualifications: {missing}")
-      
-      # Additional validation can be added here
+        # Try to find JSON without the markdown code block
+        json_pattern = r"({[\s\S]*})"
+        match = re.search(json_pattern, text)
+        if match:
+            return match.group(1).strip()
 
-  def validate_skills_mapping(self, skills_mapping: Dict) -> None:
-      """Validate the structure of skills mapping."""
-      required_keys = ["skills_mapping", "resume_recommendations"]
-      
-      if not all(key in skills_mapping for key in required_keys):
-          missing = [key for key in required_keys if key not in skills_mapping]
-          raise ValueError(f"Missing required keys in skills mapping: {missing}")
-      
-      # Additional validation can be added here
+        raise ValueError("No JSON found in the response")
 
-  def extract_header_from_resume(self, resume: str) -> str:
-      """Extract header information from the original resume."""
-      # Simple implementation - in practice, you might want a more sophisticated approach
-      lines = resume.strip().split("\n")
-      # Take the first few non-empty lines as the header
-      header_lines = []
-      for line in lines[:10]:  # Arbitrarily choose first 10 lines to search
-          if line.strip():
-              header_lines.append(line)
-          if len(header_lines) >= 4:  # Assume header is no more than 4 lines
-              break
-      return "\n".join(header_lines)
+    def validate_job_requirements(self, job_requirements: Dict) -> None:
+        """Validate the structure of job requirements."""
+        required_keys = ["job_requirements"]
 
-  # Stage 1A: Job Analysis Implementation
-  def create_job_analysis_prompt(self, job_description: str) -> tuple[str, str]:
-      """Create the prompt for job requirements analysis."""
-      logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
-      system_prompt = """
+        if not all(key in job_requirements for key in required_keys):
+            missing = [key for key in required_keys if key not in job_requirements]
+            raise ValueError(f"Missing required keys in job requirements: {missing}")
+
+        # Additional validation can be added here
+
+    def validate_candidate_qualifications(self, candidate_qualifications: Dict) -> None:
+        """Validate the structure of candidate qualifications."""
+        required_keys = ["candidate_qualifications"]
+
+        if not all(key in candidate_qualifications for key in required_keys):
+            missing = [
+                key for key in required_keys if key not in candidate_qualifications
+            ]
+            raise ValueError(
+                f"Missing required keys in candidate qualifications: {missing}"
+            )
+
+        # Additional validation can be added here
+
+    def validate_skills_mapping(self, skills_mapping: Dict) -> None:
+        """Validate the structure of skills mapping."""
+        required_keys = ["skills_mapping", "resume_recommendations"]
+
+        if not all(key in skills_mapping for key in required_keys):
+            missing = [key for key in required_keys if key not in skills_mapping]
+            raise ValueError(f"Missing required keys in skills mapping: {missing}")
+
+        # Additional validation can be added here
+
+    def extract_header_from_resume(self, resume: str) -> str:
+        """Extract header information from the original resume."""
+        # Simple implementation - in practice, you might want a more sophisticated approach
+        lines = resume.strip().split("\n")
+        # Take the first few non-empty lines as the header
+        header_lines = []
+        for line in lines[:10]:  # Arbitrarily choose first 10 lines to search
+            if line.strip():
+                header_lines.append(line)
+            if len(header_lines) >= 4:  # Assume header is no more than 4 lines
+                break
+        return "\n".join(header_lines)
+
+    def generate_resume_from_skill_assessments(
+        self,
+        candidate_name,
+        candidate_contact_info,
+        skill_assessment_results,
+        original_resume,
+    ):
+        """
+        Generate a professional resume based on skill assessment results
+
+        Parameters:
+        - candidate_name (str): The candidate's full name
+        - candidate_contact_info (dict): Contact details like email, phone, location
+        - skill_assessment_results (list): List of individual skill assessment results from LLM queries
+        - original_resume (str): Original resume text for reference
+
+        Returns:
+        - str: System prompt for generating a professional resume
+        """
+
+        # Extract and organize skill assessment data
+        skills_by_strength = {"STRONG": [], "MODERATE": [], "WEAK": [], "NONE": []}
+
+        experience_evidence = {}
+
+        # Process each skill assessment
+        for result in skill_assessment_results:
+            # Parse the result if it's a string
+            if isinstance(result, str):
+                try:
+                    import json
+
+                    result = json.loads(result)
+                    assessment = result.get("skill_assessment", {})
+                except:
+                    continue
+            else:
+                assessment = result.get("skill_assessment", {})
+
+            skill = assessment.get("skill", "")
+            strength = assessment.get("evidence_strength", "NONE")
+
+            # Add to appropriate strength category
+            if skill and strength in skills_by_strength:
+                skills_by_strength[strength].append(skill)
+
+            # Collect experience evidence
+            for evidence in assessment.get("evidence_details", []):
+                source = evidence.get("source", "")
+                if source:
+                    if source not in experience_evidence:
+                        experience_evidence[source] = []
+
+                    experience_evidence[source].append(
+                        {
+                            "skill": skill,
+                            "quote": evidence.get("quote", ""),
+                            "context": evidence.get("context", ""),
+                        }
+                    )
+
+        # Format contact info
+        contact_info_str = ""
+        if candidate_contact_info:
+            contact_items = []
+            for key, value in candidate_contact_info.items():
+                if value:
+                    contact_items.append(f"{key}: {value}")
+            contact_info_str = "\n".join(contact_items)
+
+        # Build the system prompt
+        system_prompt = f"""You are a professional resume writer with expertise in highlighting candidate strengths and experiences. 
+Create a polished, concise, and ATS-friendly resume for the candidate based on the assessment data provided.
+
+## CANDIDATE INFORMATION:
+Name: {candidate_name}
+{contact_info_str}
+
+## SKILL ASSESSMENT RESULTS:
+"""
+
+        if len(skills_by_strength["STRONG"]):
+            system_prompt += f"""\
+
+### Strong Skills (prominent in resume):
+{", ".join(skills_by_strength["STRONG"])}
+"""
+
+        if len(skills_by_strength["MODERATE"]):
+            system_prompt += f"""\
+
+### Moderate Skills (demonstrated in resume):
+{", ".join(skills_by_strength["MODERATE"])}
+"""
+
+        if len(skills_by_strength["WEAK"]):
+            system_prompt += f"""\
+
+### Weaker Skills (mentioned or implied):
+{", ".join(skills_by_strength["WEAK"])}
+"""
+            system_prompt += """\
+       
+## EXPERIENCE EVIDENCE:
+"""
+
+        # Add experience evidence by source/position
+        for source, evidences in experience_evidence.items():
+            system_prompt += f"\n### {source}:\n"
+            for evidence in evidences:
+                system_prompt += f"- {evidence['skill']}: {evidence['context']}\n"
+
+        # Add instructions for the resume creation
+        system_prompt += """\
+    
+## INSTRUCTIONS:
+
+1. Create a professional resume that emphasizes the candidate's strongest skills and most relevant experiences.
+2. Format the resume in a clean, concise, and modern style that will pass ATS systems.
+3. Include these sections:
+  - Professional Summary (highlight strongest skills and experience level)
+  - Skills (organized by strength)
+  - Professional Experience (focus on achievements and evidence of the skill)
+  - Education
+  - Additional sections as appropriate
+4. Use action verbs and quantifiable achievements where possible.
+5. Maintain a professional tone throughout.
+6. Be concise and impactful - the resume should be 1-2 pages MAXIMUM.
+7. Ensure all information is accurate to the original resume - do not embellish or fabricate experiences.
+
+## OUTPUT FORMAT:
+Provide the resume in clean markdown format, ready for the candidate to use.
+
+## REFERENCE (Original Resume):
+"""
+
+        # Add a truncated version of the original resume for reference if it's too long
+        max_resume_length = 25000  # Characters
+        if len(original_resume) > max_resume_length:
+            system_prompt += (
+                original_resume[:max_resume_length]
+                + "...\n[Original resume truncated due to length]"
+            )
+        else:
+            system_prompt += original_resume
+
+        prompt = "Create a tailored professional resume that highlights candidate's skills and experience most relevant to the job requirements. Format it in clean, ATS-friendly markdown. Provide ONLY the resume with no commentary before or after."
+        return system_prompt, prompt
+
+    async def process_and_generate_resume(
+        self,
+        message: Message,
+        job_requirements,
+        skill_assessment_results,
+        candidate_info,
+        metadata: Dict[str, Any],
+    ) -> AsyncGenerator[Message, None]:
+        """
+        Process skill assessment results and generate an optimized resume
+
+        Parameters:
+        - job_requirements (dict): The original job requirements
+        - skill_assessment_results (dict): Results from skill assessments
+        - candidate_info (dict): Basic candidate information
+
+        Returns:
+        - dict: Contains the generated resume and match analysis
+        """
+
+        # Extract candidate information
+        candidate_name = candidate_info.get("name", "")
+        candidate_contact = candidate_info.get("contact_info", {})
+        original_resume = candidate_info.get("original_resume", "")
+
+        # Generate resume prompt
+        system_prompt, prompt = self.generate_resume_from_skill_assessments(
+            candidate_name,
+            candidate_contact,
+            skill_assessment_results.values(),
+            original_resume,
+        )
+        metadata["system_prompt"] = system_prompt
+        metadata["prompt"] = prompt
+
+        message.status = "thinking"
+        message.response = "Generating tailored resume..."
+        yield message
+
+        # This would be replaced with actual LLM query
+        async for message in self.call_llm(
+            message=message, system_prompt=system_prompt, prompt=prompt
+        ):
+            if message.status != "done":
+                yield message
+        if message.status == "error":
+            return
+
+        metadata["resume"] = message.response
+        metadata["match_stats"] = self.calculate_match_statistics(
+            job_requirements, skill_assessment_results
+        )
+        # Calculate skill match statistics
+        yield message
+        return
+
+    def calculate_match_statistics(self, job_requirements, skill_assessment_results):
+        """
+        Calculate statistics about how well the candidate matches job requirements
+
+        Parameters:
+        - job_requirements (dict): The job requirements dictionary
+        - skill_assessment_results (dict): Results from skill assessments
+
+        Returns:
+        - dict: Match statistics
+        """
+
+        # Initialize counters
+        required_skills_total = 0
+        required_skills_matched = 0
+        preferred_skills_total = 0
+        preferred_skills_matched = 0
+
+        # Count required technical skills
+        tech_required = job_requirements.get("technical_skills", {}).get("required", [])
+        required_skills_total += len(tech_required)
+
+        # Count preferred technical skills
+        tech_preferred = job_requirements.get("technical_skills", {}).get(
+            "preferred", []
+        )
+        preferred_skills_total += len(tech_preferred)
+
+        # Count required experience
+        exp_required = job_requirements.get("experience_requirements", {}).get(
+            "required", []
+        )
+        required_skills_total += len(exp_required)
+
+        # Count preferred experience
+        exp_preferred = job_requirements.get("experience_requirements", {}).get(
+            "preferred", []
+        )
+        preferred_skills_total += len(exp_preferred)
+
+        # Education requirements count toward required
+        edu_required = job_requirements.get("education_requirements", [])
+        required_skills_total += len(edu_required)
+
+        # Soft skills count toward preferred
+        soft_skills = job_requirements.get("soft_skills", [])
+        preferred_skills_total += len(soft_skills)
+
+        # Industry knowledge counts toward preferred
+        industry_knowledge = job_requirements.get("industry_knowledge", [])
+        preferred_skills_total += len(industry_knowledge)
+
+        # Check matches in assessment results
+        for key, result in skill_assessment_results.items():
+            # Parse if string
+            if isinstance(result, str):
+                try:
+                    import json
+
+                    result = json.loads(result)
+                except:
+                    continue
+
+            assessment = result.get("skill_assessment", {})
+            evidence_found = assessment.get("evidence_found", False)
+            evidence_strength = assessment.get("evidence_strength", "NONE")
+
+            # Consider STRONG and MODERATE evidence as matches
+            is_match = evidence_found and evidence_strength in ["STRONG", "MODERATE"]
+
+            if not is_match:
+                continue
+
+            # Determine which category this skill belongs to
+            if key.startswith("technical_required_"):
+                required_skills_matched += 1
+            elif key.startswith("technical_preferred_"):
+                preferred_skills_matched += 1
+            elif key.startswith("experience_required_"):
+                required_skills_matched += 1
+            elif key.startswith("experience_preferred_"):
+                preferred_skills_matched += 1
+            elif key.startswith("education_"):
+                required_skills_matched += 1
+            elif key.startswith("soft_skill_"):
+                preferred_skills_matched += 1
+            elif key.startswith("industry_knowledge_"):
+                preferred_skills_matched += 1
+
+        # Calculate percentages
+        required_match_percent = (
+            (required_skills_matched / required_skills_total * 100)
+            if required_skills_total > 0
+            else 0
+        )
+        preferred_match_percent = (
+            (preferred_skills_matched / preferred_skills_total * 100)
+            if preferred_skills_total > 0
+            else 0
+        )
+        overall_total = required_skills_total + preferred_skills_total
+        overall_matched = required_skills_matched + preferred_skills_matched
+        overall_match_percent = (
+            (overall_matched / overall_total * 100) if overall_total > 0 else 0
+        )
+
+        return {
+            "required_skills": {
+                "total": required_skills_total,
+                "matched": required_skills_matched,
+                "percentage": round(required_match_percent, 1),
+            },
+            "preferred_skills": {
+                "total": preferred_skills_total,
+                "matched": preferred_skills_matched,
+                "percentage": round(preferred_match_percent, 1),
+            },
+            "overall_match": {
+                "total": overall_total,
+                "matched": overall_matched,
+                "percentage": round(overall_match_percent, 1),
+            },
+        }
+
+    def generate_skill_assessment_prompt(self, skill, resume, rag_content):
+        """
+        Generate a system prompt to query the LLM for evidence of a specific skill
+
+        Parameters:
+        - skill (str): The specific skill to assess from job requirements
+        - resume (str): The full text of the candidate's resume
+        - rag_content (str, optional): Additional RAG content queried from candidate documents
+
+        Returns:
+        - str: A system prompt tailored to assess the specific skill
+        """
+
+        system_prompt = f"""You are an objective skill assessor. Your task is to determine if a candidate possesses 
+a SPECIFIC skill based solely on their resume and supporting evidence.
+
+## SKILL TO ASSESS:
+"{skill}"
+
+## INSTRUCTIONS:
+1. Focus exclusively on assessing the candidate's proficiency with the skill: "{skill}".
+2. Examine the resume and supporting documents for both explicit mentions and clearly demonstrated applications of the skill.
+3. Do NOT infer the skill unless it is either:
+   - Directly mentioned in association with experience, or
+   - Clearly evidenced through relevant tools, technologies, responsibilities, or outcomes.
+   - Referenced as being used or listed as part of a broader initiative.
+4. Evaluate each piece of evidence and assign a confidence rating:
+   - STRONG: Explicit use with clear context or repeated/prolonged experience
+   - MODERATE: Inferred through tools, environments, or outcomes (e.g., Python used in a listed project/tool)
+   - WEAK: Mentioned in a list or indirectly implied without context
+   - NONE: No relevant evidence
+5. Provide up to 10 evidence_details. Each should include:
+   - source: where the evidence appears
+   - quote: an exact snippet from the document(s)
+   - context: a short rationale explaining how this supports the skill claim
+6. When no supporting evidence exists, output a "NONE" rating with an empty evidence details list.
+7. IMPORTANT: Even if the skill is only listed in a technologies or languages section, consider this valid evidence.
+
+## OUTPUT FORMAT:
+```json
+{{
+  "skill_assessment": {{
+    "skill": "{skill}",
+    "evidence_found": true/false,
+    "evidence_strength": "STRONG/MODERATE/WEAK/NONE",
+    "description": "short (two to three sentence) description of what {skill} means with a concise example of what you're looking for",
+    "evidence_details": [
+      {{
+        "source": "resume section/position/project",
+        "quote": "exact text from resume showing evidence",
+        "context": "brief explanation of how this demonstrates the skill"
+      }}
+    ]
+  }}
+}}
+```
+
+## CANDIDATE RESUME:
+{resume}
+"""
+
+        # Add RAG content if provided
+        if rag_content:
+            system_prompt += f"""
+
+## ADDITIONAL SUPPORTING EVIDENCE:
+{rag_content}
+"""
+
+        # Add final instruction to prevent hallucination
+        system_prompt += f"""
+
+IMPORTANT: Be factual and precise. If you cannot find strong evidence for this specific skill, it's better to indicate "evidence_strength": WEAK than to stretch for connections. Focus only on "{skill}" and not similar skills unless they directly demonstrate the required skill. Remember that a skill listed in a "Languages" or "Technologies" section should be considered valid evidence.
+"""
+
+        prompt = f'Analyze candidate resume and supporting evidence to determine if they have experience with "{skill}". Adhere strictly to the output format requested.'
+
+        return system_prompt, prompt
+
+    def process_job_requirements(self, job_requirements, resume, rag_function):
+        """
+        Process all job requirements and build individual skill assessment prompts
+
+        Parameters:
+        - job_requirements (dict): The dict returned from job analysis
+        - resume (str): The candidate's full resume
+        - rag_function (callable): Function to retrieve relevant RAG content for a skill
+
+        Returns:
+        - dict: A dictionary of skill assessment prompts
+        """
+        skill_assessment_prompts = {}
+
+        # Process technical skills - required
+        for skill in job_requirements.get("technical_skills", {}).get("required", []):
+            rag_content, metadata = rag_function(skill)
+            system_prompt, prompt = self.generate_skill_assessment_prompt(
+                skill, resume, rag_content
+            )
+            skill_assessment_prompts[f"technical_required_{skill}"] = (
+                system_prompt,
+                prompt,
+                "technical_required",
+                skill,
+                metadata,
+            )
+
+        # Process technical skills - preferred
+        for skill in job_requirements.get("technical_skills", {}).get("preferred", []):
+            rag_content, metadata = rag_function(skill)
+            system_prompt, prompt = self.generate_skill_assessment_prompt(
+                skill, resume, rag_content
+            )
+            skill_assessment_prompts[f"technical_preferred_{skill}"] = (
+                system_prompt,
+                prompt,
+                "technical_preferred",
+                skill,
+                metadata,
+            )
+
+        # Process experience requirements - required
+        for exp in job_requirements.get("experience_requirements", {}).get(
+            "required", []
+        ):
+            rag_content, metadata = rag_function(exp)
+            system_prompt, prompt = self.generate_skill_assessment_prompt(
+                exp, resume, rag_content
+            )
+            skill_assessment_prompts[f"experience_required_{exp}"] = (
+                system_prompt,
+                prompt,
+                "experience_resuired",
+                exp,
+                metadata,
+            )
+
+        # Process experience requirements - preferred
+        for exp in job_requirements.get("experience_requirements", {}).get(
+            "preferred", []
+        ):
+            rag_content, metadata = rag_function(exp)
+            system_prompt, prompt = self.generate_skill_assessment_prompt(
+                exp, resume, rag_content
+            )
+            skill_assessment_prompts[f"experience_preferred_{exp}"] = (
+                system_prompt,
+                prompt,
+                "experience_preferred",
+                exp,
+                metadata,
+            )
+
+        # Process education requirements
+        for edu in job_requirements.get("education_requirements", []):
+            rag_content, metadata = rag_function(edu)
+            system_prompt, prompt = self.generate_skill_assessment_prompt(
+                edu, resume, rag_content
+            )
+            skill_assessment_prompts[f"education_{edu}"] = (
+                system_prompt,
+                prompt,
+                "education",
+                edu,
+                metadata,
+            )
+
+        # Process soft skills
+        for skill in job_requirements.get("soft_skills", []):
+            rag_content, metadata = rag_function(skill)
+            system_prompt, prompt = self.generate_skill_assessment_prompt(
+                skill, resume, rag_content
+            )
+            skill_assessment_prompts[f"soft_skill_{skill}"] = (
+                system_prompt,
+                prompt,
+                "soft_skills",
+                skill,
+                metadata,
+            )
+
+        # Process industry knowledge
+        for knowledge in job_requirements.get("industry_knowledge", []):
+            rag_content, metadata = rag_function(knowledge)
+            system_prompt, prompt = self.generate_skill_assessment_prompt(
+                knowledge, resume, rag_content
+            )
+            skill_assessment_prompts[f"industry_knowledge_{knowledge}"] = (
+                system_prompt,
+                prompt,
+                "industry_knowledge",
+                knowledge,
+                metadata,
+            )
+
+        return skill_assessment_prompts
+
+    # Stage 1A: Job Analysis Implementation
+    def create_job_analysis_prompt(self, job_description: str) -> tuple[str, str]:
+        """Create the prompt for job requirements analysis."""
+        logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
+        system_prompt = """
       You are an objective job requirements analyzer. Your task is to extract and categorize the specific skills,
       experiences, and qualifications required in a job description WITHOUT any reference to any candidate.
       
@@ -386,76 +817,87 @@ class JobDescription(Agent):
       
       Avoid vague categorizations and be precise about whether skills are explicitly required or just preferred.
       """
-      
-      prompt = f"Job Description:\n{job_description}"
-      return system_prompt, prompt
 
-  async def analyze_job_requirements(self, message, job_description: str, metadata: Dict[str, Any]) -> AsyncGenerator[Message, None]:
-      """Analyze job requirements from job description."""
-      logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
-      try:
-          system_prompt, prompt = self.create_job_analysis_prompt(job_description)
-          metadata["system_prompt"] = system_prompt
-          metadata["prompt"] = prompt
-          async for message in self.call_llm(message=message, system_prompt=system_prompt, prompt=prompt):
-            if message.status != "done":
-              yield message
-          if message.status == "error":
-             return
-          
-          # Extract JSON from response
-          json_str = self.extract_json_from_text(message.response)
-          job_requirements = json.loads(json_str)
-          
-          self.validate_job_requirements(job_requirements)
+        prompt = f"Job Description:\n{job_description}"
+        return system_prompt, prompt
 
-          metadata["results"] = job_requirements["job_requirements"]
-          
-          message.status = "done"
-          message.response = json_str
-          yield message
-          return
+    async def analyze_job_requirements(
+        self, message, job_description: str, metadata: Dict[str, Any]
+    ) -> AsyncGenerator[Message, None]:
+        """Analyze job requirements from job description."""
+        logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
+        try:
+            system_prompt, prompt = self.create_job_analysis_prompt(job_description)
+            metadata["system_prompt"] = system_prompt
+            metadata["prompt"] = prompt
+            async for message in self.call_llm(
+                message=message, system_prompt=system_prompt, prompt=prompt
+            ):
+                if message.status != "done":
+                    yield message
+            if message.status == "error":
+                return
 
-      except Exception as e:
-          metadata["error"] = message.response
-          raise
+            # Extract JSON from response
+            json_str = self.extract_json_from_text(message.response)
+            job_requirements = json.loads(json_str)
 
-  def format_rag_context(self, rag_results: List[Dict[str, Any]]) -> str:
-      """
-      Format RAG results from process_job_requirements into a structured string.
-      
-      Args:
-          rag_results: List of dictionaries from process_job_requirements.
-      
-      Returns:
-          A formatted string for inclusion in the prompt.
-      """
-      if not rag_results:
-          return "No additional context available."
-      
-      # Group results by category and subcategory
-      grouped_context = defaultdict(list)
-      for result in rag_results:
-        key = f"{result['category']}/{result['subcategory']}".strip("/")
-        grouped_context[key].append({
-          "query": result["context"],
-          "content": result["content"][:100] + "..." if len(result["content"]) > 100 else result["content"]
-        })
-      
-      # Format as a structured string
-      context_lines = ["Additional Context from Document Retrieval:"]
-      for category, items in grouped_context.items():
-          context_lines.append(f"\nCategory: {category}")
-          for item in items:
-              context_lines.append(f"- Query: {item['query']}")
-              context_lines.append(f"  Relevant Document: {item['content']}")
-      
-      return "\n".join(context_lines)
-  
-  # Stage 1B: Candidate Analysis Implementation
-  def create_candidate_analysis_prompt(self, resume: str, rag_results: List[Dict[str, Any]]) -> tuple[str, str]:
-    """Create the prompt for candidate qualifications analysis."""
-    system_prompt = """\
+            self.validate_job_requirements(job_requirements)
+
+            metadata["results"] = job_requirements["job_requirements"]
+
+            message.status = "done"
+            yield message
+            return
+
+        except Exception as e:
+            metadata["error"] = message.response
+            raise
+
+    def format_rag_context(self, rag_results: List[Dict[str, Any]]) -> str:
+        """
+        Format RAG results from process_job_requirements into a structured string.
+
+        Args:
+            rag_results: List of dictionaries from process_job_requirements.
+
+        Returns:
+            A formatted string for inclusion in the prompt.
+        """
+        if not rag_results:
+            return "No additional context available."
+
+        # Group results by category and subcategory
+        grouped_context = defaultdict(list)
+        for result in rag_results:
+            key = f"{result['category']}/{result['subcategory']}".strip("/")
+            grouped_context[key].append(
+                {
+                    "query": result["context"],
+                    "content": (
+                        result["content"][:500] + "..."
+                        if len(result["content"]) > 500
+                        else result["content"]
+                    ),
+                }
+            )
+
+        # Format as a structured string
+        context_lines = ["Additional Context from Document Retrieval:"]
+        for category, items in grouped_context.items():
+            context_lines.append(f"\nCategory: {category}")
+            for item in items:
+                context_lines.append(f"- Query: {item['query']}")
+                context_lines.append(f"  Relevant Document: {item['content']}")
+
+        return "\n".join(context_lines)
+
+    # Stage 1B: Candidate Analysis Implementation
+    def create_candidate_analysis_prompt(
+        self, resume: str, rag_results: List[Dict[str, Any]]
+    ) -> tuple[str, str]:
+        """Create the prompt for candidate qualifications analysis."""
+        system_prompt = """\
 You are an objective resume analyzer. Create a concise inventory of the candidate's key skills, experiences, and qualifications based on their resume.
 
 CORE PRINCIPLES:
@@ -515,98 +957,110 @@ OUTPUT FORMAT:
 
 IMPORTANT: If at any point you find yourself repeating items or getting stuck, STOP that section and move to the next. It's better to provide a partial analysis than to get stuck in a loop.
 """
-    context = self.format_rag_context(rag_results)
-    prompt = f"Resume:\n{resume}\n\nAdditional Context:\n{context}"
-    return system_prompt, prompt
+        context = self.format_rag_context(rag_results)
+        prompt = f"Resume:\n{resume}\n\nAdditional Context:\n{context}"
+        return system_prompt, prompt
 
-  async def call_llm(self, message: Message, system_prompt, prompt, temperature=0.7):
-    logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
+    async def call_llm(self, message: Message, system_prompt, prompt, temperature=0.7):
+        logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
 
-    messages : List[LLMMessage] = [
-      LLMMessage(role="system", content=system_prompt),
-      LLMMessage(role="user", content=prompt)
-    ]
-    message.metadata["options"]={
-        "seed": 8911,
-        "num_ctx": self.context_size,
-        "temperature": temperature # Higher temperature to encourage tool usage
-    }
+        messages: List[LLMMessage] = [
+            LLMMessage(role="system", content=system_prompt),
+            LLMMessage(role="user", content=prompt),
+        ]
+        message.metadata["options"] = {
+            "seed": 8911,
+            "num_ctx": self.context_size,
+            "temperature": temperature,  # Higher temperature to encourage tool usage
+        }
 
-    message.status = "streaming"
-    yield message
+        message.status = "streaming"
+        yield message
 
-    last_chunk_time = 0
-    message.chunk = ""
-    message.response = ""
-    for response in self.llm.chat(
-        model=self.model, 
-        messages=messages, 
-        options={
-            **message.metadata["options"],
-        },
-        stream=True,
-    ):
-      if not response:
-          message.status = "error"
-          message.response = "No response from LLM."
-          yield message
-          return
+        last_chunk_time = 0
+        message.chunk = ""
+        message.response = ""
+        for response in self.llm.chat(
+            model=self.model,
+            messages=messages,
+            options={
+                **message.metadata["options"],
+            },
+            stream=True,
+        ):
+            if not response:
+                message.status = "error"
+                message.response = "No response from LLM."
+                yield message
+                return
 
-      message.status = "streaming"
-      message.chunk += response.message.content
-      message.response += response.message.content
+            message.status = "streaming"
+            message.chunk += response.message.content
+            message.response += response.message.content
 
-      if not response.done:
-        now = time.perf_counter()
-        if now - last_chunk_time > 0.25:
-           yield message
-           last_chunk_time = now
-           message.chunk = ""
+            if not response.done:
+                now = time.perf_counter()
+                if now - last_chunk_time > 0.25:
+                    yield message
+                    last_chunk_time = now
+                    message.chunk = ""
 
-      if response.done:
-          self.collect_metrics(response)
-          message.metadata["eval_count"] += response.eval_count
-          message.metadata["eval_duration"] += response.eval_duration
-          message.metadata["prompt_eval_count"] += response.prompt_eval_count
-          message.metadata["prompt_eval_duration"] += response.prompt_eval_duration
-          self.context_tokens = response.prompt_eval_count + response.eval_count
-          message.chunk = ""
-          message.status = "done"
-          yield message
+            if response.done:
+                self.collect_metrics(response)
+                message.metadata["eval_count"] += response.eval_count
+                message.metadata["eval_duration"] += response.eval_duration
+                message.metadata["prompt_eval_count"] += response.prompt_eval_count
+                message.metadata[
+                    "prompt_eval_duration"
+                ] += response.prompt_eval_duration
+                self.context_tokens = response.prompt_eval_count + response.eval_count
+                message.chunk = ""
+                message.status = "done"
+                yield message
 
-  async def analyze_candidate_qualifications(self, message: Message, resume: str, rag_context: List[Dict[str, Any]], metadata: Dict[str, Any]) -> AsyncGenerator[Message, None]:
-      """Analyze candidate qualifications from resume and context."""
-      try:
-          system_prompt, prompt = self.create_candidate_analysis_prompt(resume, rag_context)
-          metadata["system_prompt"] = system_prompt
-          metadata["prompt"] = prompt
-          async for message in self.call_llm(message, system_prompt, prompt):
-            if message.status != "done":
-              yield message
-          if message.status == "error":
-             return
-          
-          # Extract JSON from response
-          json_str = self.extract_json_from_text(message.response)
-          candidate_qualifications = json.loads(json_str)
-          
-          # Validate structure
-          self.validate_candidate_qualifications(candidate_qualifications)
+    async def analyze_candidate_qualifications(
+        self,
+        message: Message,
+        resume: str,
+        rag_context: List[Dict[str, Any]],
+        metadata: Dict[str, Any],
+    ) -> AsyncGenerator[Message, None]:
+        """Analyze candidate qualifications from resume and context."""
+        try:
+            system_prompt, prompt = self.create_candidate_analysis_prompt(
+                resume, rag_context
+            )
+            metadata["system_prompt"] = system_prompt
+            metadata["prompt"] = prompt
+            async for message in self.call_llm(message, system_prompt, prompt):
+                if message.status != "done":
+                    yield message
+            if message.status == "error":
+                return
 
-          metadata["results"] = candidate_qualifications["candidate_qualifications"]
-          message.status = "done"
-          message.response = json.dumps(candidate_qualifications)
-          return
-      
-      except Exception as e:
-          metadata["error"] = message.response
-          yield message
-          raise
+            # Extract JSON from response
+            json_str = self.extract_json_from_text(message.response)
+            candidate_qualifications = json.loads(json_str)
 
-  # Stage 1C: Mapping Analysis Implementation
-  def create_mapping_analysis_prompt(self ,job_requirements: Dict, candidate_qualifications: Dict) -> tuple[str, str]:
-      """Create the prompt for mapping analysis."""
-      system_prompt = """
+            # Validate structure
+            self.validate_candidate_qualifications(candidate_qualifications)
+
+            metadata["results"] = candidate_qualifications["candidate_qualifications"]
+            message.status = "done"
+            message.response = json.dumps(candidate_qualifications)
+            return
+
+        except Exception as e:
+            metadata["error"] = message.response
+            yield message
+            raise
+
+    # Stage 1C: Mapping Analysis Implementation
+    def create_mapping_analysis_prompt(
+        self, job_requirements: Dict, candidate_qualifications: Dict
+    ) -> tuple[str, str]:
+        """Create the prompt for mapping analysis."""
+        system_prompt = """
       You are an objective skills mapper. Your task is to identify legitimate matches between job requirements
       and candidate qualifications WITHOUT fabricating or stretching the truth.
       
@@ -671,48 +1125,58 @@ IMPORTANT: If at any point you find yourself repeating items or getting stuck, S
       4. Every match or transfer claim must cite specific evidence from the candidate materials
       5. Be conservative in claiming transferability - when in doubt, list as missing
       """
-      
-      prompt = f"Job Requirements:\n{json.dumps(job_requirements, indent=2)}\n\n"
-      prompt += f"Candidate Qualifications:\n{json.dumps(candidate_qualifications, indent=2)}"
-      return system_prompt, prompt
 
-  async def create_skills_mapping(self, message, job_requirements: Dict, candidate_qualifications: Dict, metadata: Dict[str, Any]) -> AsyncGenerator[Message, None]:
-      """Create mapping between job requirements and candidate qualifications."""
-      json_str = ""
-      try:
-          system_prompt, prompt = self.create_mapping_analysis_prompt(job_requirements, candidate_qualifications)
-          metadata["system_prompt"] = system_prompt
-          metadata["prompt"] = prompt
+        prompt = f"Job Requirements:\n{json.dumps(job_requirements, indent=2)}\n\n"
+        prompt += f"Candidate Qualifications:\n{json.dumps(candidate_qualifications, indent=2)}"
+        return system_prompt, prompt
 
-          async for message in self.call_llm(message, system_prompt, prompt):
-            if message != "done":
-              yield message
-          if message.status == "error":
-             return
-          
-          # Extract JSON from response
-          json_str = self.extract_json_from_text(message.response)
-          skills_mapping = json.loads(json_str)
-          
-          # Validate structure
-          self.validate_skills_mapping(skills_mapping)
+    async def create_skills_mapping(
+        self,
+        message,
+        job_requirements: Dict,
+        candidate_qualifications: Dict,
+        metadata: Dict[str, Any],
+    ) -> AsyncGenerator[Message, None]:
+        """Create mapping between job requirements and candidate qualifications."""
+        json_str = ""
+        try:
+            system_prompt, prompt = self.create_mapping_analysis_prompt(
+                job_requirements, candidate_qualifications
+            )
+            metadata["system_prompt"] = system_prompt
+            metadata["prompt"] = prompt
 
-          metadata["skills_mapping"] = skills_mapping["skills_mapping"]
-          
-          message.status = "done"
-          message.response = json_str
-          yield message
-          return
-      
-      except Exception as e:
-          metadata["error"] = json_str
-          yield message
-          raise
+            async for message in self.call_llm(message, system_prompt, prompt):
+                if message != "done":
+                    yield message
+            if message.status == "error":
+                return
 
-  # Stage 2: Resume Generation Implementation
-  def create_resume_generation_prompt(self, skills_mapping: Dict, candidate_qualifications: Dict, original_header: str) -> tuple[str, str]:
-      """Create the prompt for resume generation."""
-      system_prompt = """
+            # Extract JSON from response
+            json_str = self.extract_json_from_text(message.response)
+            skills_mapping = json.loads(json_str)
+
+            # Validate structure
+            self.validate_skills_mapping(skills_mapping)
+
+            metadata["skills_mapping"] = skills_mapping["skills_mapping"]
+
+            message.status = "done"
+            message.response = json_str
+            yield message
+            return
+
+        except Exception as e:
+            metadata["error"] = json_str
+            yield message
+            raise
+
+    # Stage 2: Resume Generation Implementation
+    def create_resume_generation_prompt(
+        self, skills_mapping: Dict, candidate_qualifications: Dict, original_header: str
+    ) -> tuple[str, str]:
+        """Create the prompt for resume generation."""
+        system_prompt = """
       You are a professional resume writer whose primary concern is FACTUAL ACCURACY. Your task is to create
       a tailored resume that presents the candidate's actual qualifications in the most relevant way for this job,
       using ONLY information that has been verified in the skills mapping.
@@ -759,30 +1223,46 @@ IMPORTANT: If at any point you find yourself repeating items or getting stuck, S
       3. Ensure all experience descriptions can be traced to evidence in candidate qualifications
       4. Confirm that transferable skills are presented honestly without exaggeration
       """
-      
-      prompt = f"Skills Mapping:\n{json.dumps(skills_mapping, indent=2)}\n\n"
-      prompt += f"Candidate Qualifications:\n{json.dumps(candidate_qualifications, indent=2)}\n\n"
-      prompt += f"Original Resume Header:\n{original_header}"
-      return system_prompt, prompt
 
-  async def generate_tailored_resume(self, message, skills_mapping: Dict, candidate_qualifications: Dict, original_header: str, metadata: Dict[str, Any]) -> AsyncGenerator[Message, None]:
-      """Generate a tailored resume based on skills mapping."""
-      system_prompt, prompt = self.create_resume_generation_prompt(skills_mapping, candidate_qualifications, original_header)
-      metadata["system_prompt"] = system_prompt
-      metadata["prompt"] = prompt
-      async for message in self.call_llm(message, system_prompt, prompt, temperature=0.4):  # Slightly higher temperature for better writing
-        if message.status != "done":
-          yield message
-      if message.status == "error":
-          return
-      metadata["results"] = message.response
-      yield message
-      return      
+        prompt = f"Skills Mapping:\n{json.dumps(skills_mapping, indent=2)}\n\n"
+        prompt += f"Candidate Qualifications:\n{json.dumps(candidate_qualifications, indent=2)}\n\n"
+        prompt += f"Original Resume Header:\n{original_header}"
+        return system_prompt, prompt
 
-  # Stage 3: Verification Implementation
-  def create_verification_prompt(self, generated_resume: str, skills_mapping: Dict, candidate_qualifications: Dict) -> tuple[str, str]:
-      """Create the prompt for resume verification."""
-      system_prompt = """
+    async def generate_tailored_resume(
+        self,
+        message,
+        skills_mapping: Dict,
+        candidate_qualifications: Dict,
+        original_header: str,
+        metadata: Dict[str, Any],
+    ) -> AsyncGenerator[Message, None]:
+        """Generate a tailored resume based on skills mapping."""
+        system_prompt, prompt = self.create_resume_generation_prompt(
+            skills_mapping, candidate_qualifications, original_header
+        )
+        metadata["system_prompt"] = system_prompt
+        metadata["prompt"] = prompt
+        async for message in self.call_llm(
+            message, system_prompt, prompt, temperature=0.4
+        ):  # Slightly higher temperature for better writing
+            if message.status != "done":
+                yield message
+        if message.status == "error":
+            return
+        metadata["results"] = message.response
+        yield message
+        return
+
+    # Stage 3: Verification Implementation
+    def create_verification_prompt(
+        self,
+        generated_resume: str,
+        skills_mapping: Dict,
+        candidate_qualifications: Dict,
+    ) -> tuple[str, str]:
+        """Create the prompt for resume verification."""
+        system_prompt = """
       You are a critical resume fact-checker responsible for verifying the accuracy of a tailored resume.
       Your task is to identify and flag any fabricated or embellished information that does not appear in
       the candidate's original materials.
@@ -855,45 +1335,66 @@ IMPORTANT: If at any point you find yourself repeating items or getting stuck, S
       - Exaggerating level of involvement or responsibility in projects or roles
       - Creating achievements that weren't documented in the original materials
       """
-      
-      prompt = f"Tailored Resume:\n{generated_resume}\n\n"
-      prompt += f"Skills Mapping:\n{json.dumps(skills_mapping, indent=2)}\n\n"
-      prompt += f"Candidate Qualifications:\n{json.dumps(candidate_qualifications, indent=2)}"
-      return system_prompt, prompt
 
-  async def verify_resume(self, message: Message, generated_resume: str, skills_mapping: Dict, candidate_qualifications: Dict, metadata: Dict[str, Any]) -> AsyncGenerator[Message, None]:
-      """Verify the generated resume for accuracy against original materials."""
-      try:
-          system_prompt, prompt = self.create_verification_prompt(generated_resume, skills_mapping, candidate_qualifications)
-          metadata["system_prompt"] = system_prompt
-          metadata["prompt"] = prompt
-          async for message in self.call_llm(message, system_prompt, prompt):
-            if message.status != "done":
-              yield message
-          if message.status == "error":
-             return
-          
-          # Extract JSON from response
-          json_str = self.extract_json_from_text(message.response)
-          metadata["results"] = json.loads(json_str)["verification_results"]
-          message.status = "done"
-          message.response = json_str
-          yield message
-          return
-      except Exception as e:
-          metadata["error"] = message.response
-          yield message
-          raise
+        prompt = f"Tailored Resume:\n{generated_resume}\n\n"
+        prompt += f"Skills Mapping:\n{json.dumps(skills_mapping, indent=2)}\n\n"
+        prompt += f"Candidate Qualifications:\n{json.dumps(candidate_qualifications, indent=2)}"
+        return system_prompt, prompt
 
-  async def correct_resume_issues(self, message: Message, generated_resume: str, verification_results: Dict, skills_mapping: Dict, candidate_qualifications: Dict, original_header: str, metadata: Dict[str, Any])  -> AsyncGenerator[Message, None]:
-      """Correct issues in the resume based on verification results."""
-      if verification_results["verification_results"]["overall_assessment"] == "APPROVED":
-        message.status = "done"
-        message.status = generated_resume
-        yield message
-        return
-      
-      system_prompt = """\
+    async def verify_resume(
+        self,
+        message: Message,
+        generated_resume: str,
+        skills_mapping: Dict,
+        candidate_qualifications: Dict,
+        metadata: Dict[str, Any],
+    ) -> AsyncGenerator[Message, None]:
+        """Verify the generated resume for accuracy against original materials."""
+        try:
+            system_prompt, prompt = self.create_verification_prompt(
+                generated_resume, skills_mapping, candidate_qualifications
+            )
+            metadata["system_prompt"] = system_prompt
+            metadata["prompt"] = prompt
+            async for message in self.call_llm(message, system_prompt, prompt):
+                if message.status != "done":
+                    yield message
+            if message.status == "error":
+                return
+
+            # Extract JSON from response
+            json_str = self.extract_json_from_text(message.response)
+            metadata["results"] = json.loads(json_str)["verification_results"]
+            message.status = "done"
+            message.response = json_str
+            yield message
+            return
+        except Exception as e:
+            metadata["error"] = message.response
+            yield message
+            raise
+
+    async def correct_resume_issues(
+        self,
+        message: Message,
+        generated_resume: str,
+        verification_results: Dict,
+        skills_mapping: Dict,
+        candidate_qualifications: Dict,
+        original_header: str,
+        metadata: Dict[str, Any],
+    ) -> AsyncGenerator[Message, None]:
+        """Correct issues in the resume based on verification results."""
+        if (
+            verification_results["verification_results"]["overall_assessment"]
+            == "APPROVED"
+        ):
+            message.status = "done"
+            message.status = generated_resume
+            yield message
+            return
+
+        system_prompt = """\
 You are a professional resume editor with a focus on factual accuracy. Your task is to correct identified issues in a tailored resume according to the verification report.
 
 ## REFERENCE DATA:
@@ -929,328 +1430,470 @@ The following sections contain reference information for you to use when making
 
 Your output should contain ONLY the corrected resume text with no additional explanations or context.
 """
-      prompt = """
+        prompt = """
 ## REFERENCE DATA
 
 ### Original Resume:
 """
-      prompt += generated_resume
-      prompt += """
+        prompt += generated_resume
+        prompt += """
 
 ### Verification Results:
 """
-      prompt += json.dumps(verification_results, indent=2)
-      prompt += """
+        prompt += json.dumps(verification_results, indent=2)
+        prompt += """
 
 ### Skills Mapping:
 """
-      prompt += json.dumps(skills_mapping, indent=2)
-      prompt += """
+        prompt += json.dumps(skills_mapping, indent=2)
+        prompt += """
 
 ### Candidate Qualifications:
 """
-      prompt += json.dumps(candidate_qualifications, indent=2)
-      prompt += """
+        prompt += json.dumps(candidate_qualifications, indent=2)
+        prompt += """
 
 ### Original Resume Header:
 """
-      prompt += generated_resume
-      prompt += """
+        prompt += generated_resume
+        prompt += """
 
 ## TASK
 Based on the reference data above, please create a corrected version of the resume that addresses all issues identified in the verification report. Return ONLY the corrected resume.
 """
-      
-      metadata["system_prompt"] = system_prompt
-      metadata["prompt"] = prompt
 
-      async for message in self.call_llm(message, prompt, system_prompt, temperature=0.3):
-        if message.status != "done":
-          yield message
-      if message.status == "error":
-          return
-      metadata["results"] = message.response
-      yield message
+        metadata["system_prompt"] = system_prompt
+        metadata["prompt"] = prompt
 
-  def process_job_requirements(self, job_requirements: Dict[str, Any]) -> List[Dict[str, Any]]:
-    """
-    Process job requirements JSON, gather RAG documents using find_similar, remove duplicates, 
-    and return top 20 results.
-    
-    Args:
-        job_requirements: Dictionary containing job requirements.
-        retriever: Instance of RagRetriever with find_similar method.
-    
-    Returns:
-        List of up to 20 ChromaDB documents, sorted by combined importance and similarity score.
-    """
-    if self.context is None or self.context.file_watcher is None:
-       raise ValueError(f"context or file_watcher is None on {self.agent_type}")
-
-    retriever = self.context.file_watcher
-    # Importance weights for each category
-    importance_weights = {
-        ("technical_skills", "required"): 1.0,
-        ("technical_skills", "preferred"): 0.8,
-        ("experience_requirements", "required"): 0.95,
-        ("experience_requirements", "preferred"): 0.75,
-        ("education_requirements", ""): 0.7,
-        ("soft_skills", ""): 0.6,
-        ("industry_knowledge", ""): 0.65,
-        ("responsibilities", ""): 0.85,
-        ("company_values", ""): 0.5
-    }
-    
-    # Store all RAG results with metadata
-    all_results = []
-    
-    def traverse_requirements(data: Any, category: str = "", subcategory: str = ""):
-        """
-        Recursively traverse the job requirements and gather RAG documents.
-        """
-        if isinstance(data, dict):
-            for key, value in data.items():
-                new_subcategory = key if category else ""
-                traverse_requirements(value, category or key, new_subcategory)
-        elif isinstance(data, list):
-            for item in data:
-                # Determine the weight key
-                weight_key = (category, subcategory) if subcategory else (category, "")
-                weight = importance_weights.get(weight_key, 0.5)  # Default weight
-                
-                # Call find_similar for the item
-                try:
-                    rag_results = retriever.find_similar(item, top_k=20, threshold=0.4) # Strict matching
-                    # Process each result
-                    for doc_id, content, distance, metadata in zip(
-                        rag_results["ids"], 
-                        rag_results["documents"], 
-                        rag_results["distances"], 
-                        rag_results["metadatas"]
-                    ):
-                        # Convert cosine distance to similarity score (higher is better)
-                        similarity_score = 1 - distance  # Cosine distance to similarity
-                        all_results.append({
-                            "id": doc_id,
-                            "content": content,
-                            "score": similarity_score,
-                            "weight": weight,
-                            "context": item,
-                            "category": category,
-                            "subcategory": subcategory,
-                            "metadata": metadata
-                        })
-                except Exception as e:
-                    logger.error(f"Error processing context '{item}': {e}")
-    
-    # Start traversal
-    traverse_requirements(job_requirements)
-    
-    # Remove duplicates based on document ID
-    unique_results = []
-    seen_ids = set()
-    for result in all_results:
-        if result["id"] not in seen_ids:
-            seen_ids.add(result["id"])
-            unique_results.append(result)
-    
-    # Sort by combined score (weight * similarity score)
-    sorted_results = sorted(
-        unique_results,
-        key=lambda x: x["weight"] * x["score"],
-        reverse=True
-    )
-    
-    # Return top 10 results
-    return sorted_results[:10]
-
-
-  async def generate_rag_content(self, message: Message, job_requirements: Dict[str, Any]) -> AsyncGenerator[Message, None]:
-    results = self.process_job_requirements(job_requirements = job_requirements)
-    message.response = f"Retrieved {len(results)} documents."
-    message.metadata["rag_context"] = results
-#     for result in results:
-#         message.response += f"""
-# ID: {result['id']}, Context: {result['context']}, \
-# Category: {result['category']}/{result['subcategory']}, \
-# Similarity Score: {result['score']:.3f}, \
-# Combined Score: {result['weight'] * result['score']:.3f}, \
-# Content: {result['content']}
-# """
-    message.status = "done"
-    yield message
-    return
-  
-  async def generate_factual_tailored_resume(self, message: Message, job_description: str, resume: str) -> AsyncGenerator[Message, None]:
-    """
-    Main function to generate a factually accurate tailored resume.
-    
-    Args:
-        job_description: The job description text
-        resume: The candidate's original resume text
-        additional_context: Any additional context about the candidate (optional)
-        
-    Returns:
-        Dict containing the generated resume and supporting analysis
-    """
-    if self.context is None:
-       raise ValueError(f"context is None in {self.agent_type}")
-    
-    message.status = "thinking"
-    logger.info(message.response)
-    yield message
-
-    message.metadata["generate_factual_tailored_resume"] = {}
-    metadata = message.metadata["generate_factual_tailored_resume"]
-    # Stage 1A: Analyze job requirements
-    message.response = "Multi-stage RAG resume generation process: Stage 1A: Analyzing job requirements"
-    logger.info(message.response)
-    yield message
-    metadata["job_requirements"] = {}
-    async for message in self.analyze_job_requirements(message, job_description, metadata["job_requirements"]):
-      if message.status != "done":
-        yield message
-    if message.status == "error":
-      return
-    job_requirements = json.loads(message.response)
-
-    # Stage 1B: Analyze candidate qualifications
-    message.status = "thinking"
-    message.response = "Multi-stage RAG resume generation process: Stage 1B: Analyzing candidate qualifications"
-    logger.info(message.response)
-    yield message
-
-    async for message in self.generate_rag_content(message, job_requirements):
-       if message.status != "done":
-        yield message
-    if message.status == "error":
-      return
-
-    rag_context = message.metadata["rag_context"]
-    metadata["analyze_candidate_qualifications"] = {
-        "rag_context": rag_context
-    }
-    async for message in self.analyze_candidate_qualifications(
-      message=message,
-      resume=resume,
-      rag_context=rag_context, 
-      metadata=metadata["analyze_candidate_qualifications"]):
-      if message.status != "done":
-        yield message
-    if message.status == "error":
-      return
-    candidate_qualifications = json.loads(message.response)
-    
-    # Stage 1C: Create skills mapping
-    message.status = "thinking"
-    message.response = "Multi-stage RAG resume generation process: Stage 1C: Creating skills mapping"
-    logger.info(message.response)
-    yield message
-    metadata["skills_mapping"] = {}
-    async for message in self.create_skills_mapping(
-        message=message, 
-        job_requirements=job_requirements, 
-        candidate_qualifications=candidate_qualifications, 
-        metadata=metadata["skills_mapping"]):
-      if message.status != "done":
-        yield message
-    if message.status == "error":
-      return
-    skills_mapping = json.loads(message.response)
-    
-    # Extract header from original resume
-    original_header = self.extract_header_from_resume(resume)
-    
-    # Stage 2: Generate tailored resume
-    message.status = "thinking"
-    message.response = "Multi-stage RAG resume generation process: Stage 2: Generating tailored resume"
-    logger.info(message.response)
-    yield message
-    metadata["generate_tailored_resume"] = {}
-    async for message in self.generate_tailored_resume(
-      message=message,
-      skills_mapping=skills_mapping,
-      candidate_qualifications=candidate_qualifications, 
-      original_header=original_header, 
-      metadata=metadata["generate_tailored_resume"]):
-      if message.status != "done":
-        yield message
-    if message.status == "error":
-      return
-    generated_resume = message.response
-
-    # Stage 3: Verify resume
-    message.status = "thinking"
-    message.response = "Multi-stage RAG resume generation process: Stage 3: Verifying resume for accuracy"
-    logger.info(message.response)
-    yield message
-    metadata["verify_resume"] = {
-        "first_pass": {}
-    }
-    async for message in self.verify_resume(
-        message=message, 
-        generated_resume=generated_resume, 
-        skills_mapping=skills_mapping, 
-        candidate_qualifications=candidate_qualifications, 
-        metadata=metadata["verify_resume"]["first_pass"]):
-      if message.status != "done":
-        yield message
-    if message.status == "error":
-      return
-    
-    verification_results = json.loads(message.response)
-    
-    # Handle corrections if needed
-    if verification_results["verification_results"]["overall_assessment"] == "NEEDS REVISION":
-        message.status = "thinking"
-        message.response = "Correcting issues found in verification"
-        logger.info(message.response)
-        yield message
-
-        metadata["correct_resume_issues"] = {}
-        async for message in self.correct_resume_issues(
-          message=message,
-          generated_resume=generated_resume, 
-          verification_results=verification_results, 
-          skills_mapping=skills_mapping, 
-          candidate_qualifications=candidate_qualifications,
-          original_header=original_header,
-          metadata=metadata["correct_resume_issues"]
+        async for message in self.call_llm(
+            message, prompt, system_prompt, temperature=0.3
         ):
-          if message.status != "done":
-            yield message
-        if message.status == "error":
-          return
-
-        generated_resume = message.response
-
-        # Re-verify after corrections
-        message.status = "thinking"
-        message.response = "Re-verifying corrected resume"
-        yield message
-
-        logger.info(message.response)
-        metadata["verify_resume"]["second_pass"] = {}
-        async for message in self.verify_resume(
-            message=message, 
-            generated_resume=generated_resume, 
-            skills_mapping=skills_mapping, 
-            candidate_qualifications=candidate_qualifications,
-            metadata=metadata["verify_resume"]["second_pass"]):
-          if message.status != "done":
-              yield message
+            if message.status != "done":
+                yield message
         if message.status == "error":
             return
+        metadata["results"] = message.response
+        yield message
+
+    def process_job_requirements2(
+        self, job_requirements: Dict[str, Any]
+    ) -> List[Dict[str, Any]]:
+        """
+        Process job requirements JSON, gather RAG documents using find_similar, remove duplicates,
+        and return top 20 results.
+
+        Args:
+            job_requirements: Dictionary containing job requirements.
+            retriever: Instance of RagRetriever with find_similar method.
+
+        Returns:
+            List of up to 20 ChromaDB documents, sorted by combined importance and similarity score.
+        """
+        if self.context is None or self.context.file_watcher is None:
+            raise ValueError(f"context or file_watcher is None on {self.agent_type}")
+
+        retriever = self.context.file_watcher
+        # Importance weights for each category
+        importance_weights = {
+            ("technical_skills", "required"): 1.0,
+            ("technical_skills", "preferred"): 0.8,
+            ("experience_requirements", "required"): 0.95,
+            ("experience_requirements", "preferred"): 0.75,
+            ("education_requirements", ""): 0.7,
+            ("soft_skills", ""): 0.6,
+            ("industry_knowledge", ""): 0.65,
+            ("responsibilities", ""): 0.85,
+            ("company_values", ""): 0.5,
+        }
+
+        # Store all RAG results with metadata
+        all_results = []
+
+        def traverse_requirements(data: Any, category: str = "", subcategory: str = ""):
+            """
+            Recursively traverse the job requirements and gather RAG documents.
+            """
+            if isinstance(data, dict):
+                for key, value in data.items():
+                    new_subcategory = key if category else ""
+                    traverse_requirements(value, category or key, new_subcategory)
+            elif isinstance(data, list):
+                for item in data:
+                    # Determine the weight key
+                    weight_key = (
+                        (category, subcategory) if subcategory else (category, "")
+                    )
+                    weight = importance_weights.get(weight_key, 0.5)  # Default weight
+
+                    # Call find_similar for the item
+                    try:
+                        rag_results = retriever.find_similar(
+                            item, top_k=20, threshold=0.5
+                        )  # Moderate matching
+                        # Process each result
+                        for doc_id, content, distance, metadata in zip(
+                            rag_results["ids"],
+                            rag_results["documents"],
+                            rag_results["distances"],
+                            rag_results["metadatas"],
+                        ):
+                            # Convert cosine distance to similarity score (higher is better)
+                            similarity_score = (
+                                1 - distance
+                            )  # Cosine distance to similarity
+                            all_results.append(
+                                {
+                                    "id": doc_id,
+                                    "content": content,
+                                    "score": similarity_score,
+                                    "weight": weight,
+                                    "context": item,
+                                    "category": category,
+                                    "subcategory": subcategory,
+                                    "metadata": metadata,
+                                }
+                            )
+                    except Exception as e:
+                        logger.error(f"Error processing context '{item}': {e}")
+
+        # Start traversal
+        traverse_requirements(job_requirements)
+
+        # Remove duplicates based on document ID
+        unique_results = []
+        seen_ids = set()
+        for result in all_results:
+            if result["id"] not in seen_ids:
+                seen_ids.add(result["id"])
+                unique_results.append(result)
+
+        # Sort by combined score (weight * similarity score)
+        sorted_results = sorted(
+            unique_results, key=lambda x: x["weight"] * x["score"], reverse=True
+        )
+
+        # Return top 10 results
+        return sorted_results[:10]
+
+    async def generate_rag_content(
+        self, message: Message, job_requirements: Dict[str, Any]
+    ) -> AsyncGenerator[Message, None]:
+        results = self.process_job_requirements2(job_requirements=job_requirements)
+        message.response = f"Retrieved {len(results)} documents."
+        message.metadata["rag_context"] = results
+        #     for result in results:
+        #         message.response += f"""
+        # ID: {result['id']}, Context: {result['context']}, \
+        # Category: {result['category']}/{result['subcategory']}, \
+        # Similarity Score: {result['score']:.3f}, \
+        # Combined Score: {result['weight'] * result['score']:.3f}, \
+        # Content: {result['content']}
+        # """
+        message.status = "done"
+        yield message
+        return
+
+    def rag_function(self, skill: str) -> tuple[str, list[Any]]:
+        if self.context is None or self.context.file_watcher is None:
+            raise ValueError("self.context or self.context.file_watcher is None")
+
+        try:
+            rag_results = ""
+            all_metadata = []
+            chroma_results = self.context.file_watcher.find_similar(
+                query=skill, top_k=5, threshold=0.5
+            )
+            if chroma_results:
+                chroma_embedding = np.array(
+                    chroma_results["query_embedding"]
+                ).flatten()  # Ensure correct shape
+                print(f"Chroma embedding shape: {chroma_embedding.shape}")
+
+                umap_2d = self.context.file_watcher.umap_model_2d.transform(
+                    [chroma_embedding]
+                )[0].tolist()
+                print(
+                    f"UMAP 2D output: {umap_2d}, length: {len(umap_2d)}"
+                )  # Debug output
+
+                umap_3d = self.context.file_watcher.umap_model_3d.transform(
+                    [chroma_embedding]
+                )[0].tolist()
+                print(
+                    f"UMAP 3D output: {umap_3d}, length: {len(umap_3d)}"
+                )  # Debug output
+
+                all_metadata.append(
+                    {
+                        "name": "JPK",
+                        **chroma_results,
+                        "umap_embedding_2d": umap_2d,
+                        "umap_embedding_3d": umap_3d,
+                    }
+                )
+
+                for index, metadata in enumerate(chroma_results["metadatas"]):
+                    content = "\n".join(
+                        [
+                            line.strip()
+                            for line in chroma_results["documents"][index].split("\n")
+                            if line
+                        ]
+                    ).strip()
+                    rag_results += f"""
+Source: {metadata.get("doc_type", "unknown")}: {metadata.get("path", "")} lines {metadata.get("line_begin", 0)}-{metadata.get("line_end", 0)}
+Content: { content }
+
+"""
+            return rag_results, all_metadata
+
+        except Exception as e:
+            logger.error(e)
+            logger.error(traceback.format_exc())
+            exit(0)
+
+    async def generate_factual_tailored_resume(
+        self, message: Message, job_description: str, resume: str
+    ) -> AsyncGenerator[Message, None]:
+        """
+        Main function to generate a factually accurate tailored resume.
+
+        Args:
+            job_description: The job description text
+            resume: The candidate's original resume text
+            additional_context: Any additional context about the candidate (optional)
+
+        Returns:
+            Dict containing the generated resume and supporting analysis
+        """
+        if self.context is None:
+            raise ValueError(f"context is None in {self.agent_type}")
+
+        message.status = "thinking"
+        logger.info(message.response)
+        yield message
+
+        message.metadata["resume_generation"] = {}
+        metadata = message.metadata["resume_generation"]
+        # Stage 1A: Analyze job requirements
+        streaming_message = Message(prompt="Analyze job requirements")
+        streaming_message.status = "thinking"
+        streaming_message.response = f"Analyzing job requirements\n\n"
+        yield streaming_message
+
+        metadata["requirements"] = {}
+        async for message in self.analyze_job_requirements(
+            message, job_description, metadata["requirements"]
+        ):
+            if message.status != "done":
+                streaming_message.response += "."
+                yield streaming_message
+        if message.status == "error":
+            return
+
+        job_requirements = metadata.get("requirements", {}).get("results", {})
+        metadata["candidate_skills"] = {}
+        candidate_skills = metadata["candidate_skills"]
+        prompts = self.process_job_requirements(
+            job_requirements=job_requirements,
+            resume=resume,
+            rag_function=self.rag_function,
+        )  # , retrieve_rag_content)
+
+        # UI should persist this state of the message
+        partial = message.model_copy()
+        partial.status = "partial"
+        partial.title = f"Found {len(prompts)} job requirements"
+        yield partial
+        self.conversation.add(partial)
+
+        # Stage 1B: Analyze candidate qualifications
+        message.status = "thinking"
+        message.response = "Performing skill-based assessment"
+        logger.info(message.response)
+        yield message
+
+        # Use each prompt to query an LLM
+        total_prompts = len(prompts.items())
+        index = 0
+        skill_assessment_results: Dict[str, Any] = {}
+        for skill_name, (
+            system_prompt,
+            prompt,
+            section,
+            description,
+            rag,
+        ) in prompts.items():
+            index += 1
+            # Query LLM with prompt
+            streaming_message = Message(prompt=prompt)
+            streaming_message.status = "thinking"
+            streaming_message.response = (
+                f"Analyzing skill {index}/{total_prompts}:\n\n* {description}\n\n"
+            )
+            yield streaming_message
+
+            async for message in self.call_llm(
+                message=message, system_prompt=system_prompt, prompt=prompt
+            ):
+                if message.status != "done":
+                    streaming_message.response += "."
+                    yield streaming_message
+            if message.status == "error":
+                return
+
+            name = re.sub(rf"^{section}_", "", skill_name)
+            json_str = self.extract_json_from_text(message.response)
+            try:
+                skill_assessment_results[skill_name] = json.loads(json_str)
+                match_level = (
+                    skill_assessment_results[skill_name]
+                    .get("skill_assessment", {})
+                    .get("evidence_strength", "UNKNOWN")
+                )
+                skill_description = (
+                    skill_assessment_results[skill_name]
+                    .get("skill_assessment", {})
+                    .get("description", "")
+                )
+                if skill_description:  # Do not duplicate this item in the JSON data
+                    del skill_assessment_results[skill_name]["skill_assessment"][
+                        "description"
+                    ]
+            except Exception as e:
+                skill_assessment_results[skill_name] = json_str
+                match_level = "UNKNOWN"
+                skill_description = ""
+
+            # UI should persist this state of the message
+            partial = message.model_copy()
+            partial.status = "partial"
+            partial.prompt = prompt
+            partial.system_prompt = system_prompt
+            partial.title = (
+                f"Skill {index}/{total_prompts}: {description} [{match_level}]"
+            )
+            partial.metadata["rag"] = rag
+            if skill_description:
+                partial.response += f"\n\n{skill_description}"
+            yield partial
+            self.conversation.add(partial)
+
+            if section not in candidate_skills:
+                candidate_skills[section] = {}
+
+            candidate_skills[section][name] = {
+                "result": skill_assessment_results[skill_name],
+            }
+
+        # Extract header from original resume
+        original_header = self.extract_header_from_resume(resume)
+        candidate_info = {
+            "name": "James Ketrenos",
+            "contact_info": {
+                "email": "james@ketrenos.com",
+                "phone": "(503) 501 8281",
+                "location": "Beaverton, OR 97003",
+            },
+            "original_resume": resume,
+        }
+
+        # Stage 2: Generate tailored resume
+        message.status = "thinking"
+        message.response = "Generating tailored resume"
+        logger.info(message.response)
+        yield message
+
+        async for message in self.process_and_generate_resume(
+            message=message,
+            job_requirements=job_requirements,
+            skill_assessment_results=skill_assessment_results,
+            candidate_info=candidate_info,
+            metadata=metadata,
+        ):
+            if message.status != "done":
+                yield message
+        if message.status == "error":
+            return
+
+        yield message
+        return
+
+        # Stage 3: Verify resume
+        message.status = "thinking"
+        message.response = "Multi-stage RAG resume generation process: Stage 3: Verifying resume for accuracy"
+        logger.info(message.response)
+        yield message
+        metadata["verify_resume"] = {"first_pass": {}}
+        async for message in self.verify_resume(
+            message=message,
+            generated_resume=generated_resume,
+            skills_mapping=skills_mapping,
+            candidate_qualifications=candidate_qualifications,
+            metadata=metadata["verify_resume"]["first_pass"],
+        ):
+            if message.status != "done":
+                yield message
+        if message.status == "error":
+            return
+
         verification_results = json.loads(message.response)
 
-    # Return the final results
-    message.status = "done"
-    message.response = generated_resume
-    yield message
+        # Handle corrections if needed
+        if (
+            verification_results["verification_results"]["overall_assessment"]
+            == "NEEDS REVISION"
+        ):
+            message.status = "thinking"
+            message.response = "Correcting issues found in verification"
+            logger.info(message.response)
+            yield message
+
+            metadata["correct_resume_issues"] = {}
+            async for message in self.correct_resume_issues(
+                message=message,
+                generated_resume=generated_resume,
+                verification_results=verification_results,
+                skills_mapping=skills_mapping,
+                candidate_qualifications=candidate_qualifications,
+                original_header=original_header,
+                metadata=metadata["correct_resume_issues"],
+            ):
+                if message.status != "done":
+                    yield message
+            if message.status == "error":
+                return
+
+            generated_resume = message.response
+
+            # Re-verify after corrections
+            message.status = "thinking"
+            message.response = "Re-verifying corrected resume"
+            yield message
+
+            logger.info(message.response)
+            metadata["verify_resume"]["second_pass"] = {}
+            async for message in self.verify_resume(
+                message=message,
+                generated_resume=generated_resume,
+                skills_mapping=skills_mapping,
+                candidate_qualifications=candidate_qualifications,
+                metadata=metadata["verify_resume"]["second_pass"],
+            ):
+                if message.status != "done":
+                    yield message
+            if message.status == "error":
+                return
+            verification_results = json.loads(message.response)
+
+        # Return the final results
+        message.status = "done"
+        message.response = generated_resume
+        yield message
+
+        logger.info("Resume generation process completed successfully")
+        return
 
-    logger.info("Resume generation process completed successfully")
-    return
-       
 
 # Register the base agent
 agent_registry.register(JobDescription._agent_type, JobDescription)
diff --git a/src/utils/agents/resume.py b/src/utils/agents/resume.py
index 3a058d4..dd7118b 100644
--- a/src/utils/agents/resume.py
+++ b/src/utils/agents/resume.py
@@ -1,12 +1,20 @@
 from __future__ import annotations
-from pydantic import model_validator # type: ignore
-from typing import Literal, ClassVar, Optional, Any, AsyncGenerator, List # NOTE: You must import Optional for late binding to work
+from pydantic import model_validator  # type: ignore
+from typing import (
+    Literal,
+    ClassVar,
+    Optional,
+    Any,
+    AsyncGenerator,
+    List,
+)  # NOTE: You must import Optional for late binding to work
 from datetime import datetime
 import inspect
 
-from . base import Agent, agent_registry
-from .. message import Message
-from .. setup_logging import setup_logging
+from .base import Agent, agent_registry
+from ..message import Message
+from ..setup_logging import setup_logging
+
 logger = setup_logging()
 
 system_fact_check = f"""
@@ -36,82 +44,94 @@ When answering queries, follow these steps:
 - Avoid phrases like 'According to the <|context|>' or similar references to the <|context|>, <|job_description|>, <|resume|>, or <|context|> tags.
 """.strip()
 
+
 class Resume(Agent):
-  agent_type: Literal["resume"] = "resume"  # type: ignore
-  _agent_type: ClassVar[str] = agent_type  # Add this for registration
+    agent_type: Literal["resume"] = "resume"  # type: ignore
+    _agent_type: ClassVar[str] = agent_type  # Add this for registration
 
-  system_prompt:str = system_fact_check
-  resume: str
+    system_prompt: str = system_fact_check
+    resume: str
 
-  @model_validator(mode="after")
-  def validate_resume(self):
-    if not self.resume.strip():
-      raise ValueError("Resume content cannot be empty")
-    return self
+    @model_validator(mode="after")
+    def validate_resume(self):
+        if not self.resume.strip():
+            raise ValueError("Resume content cannot be empty")
+        return self
 
-  async def prepare_message(self, message:Message) -> AsyncGenerator[Message, None]:
-    logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
-    if not self.context:
-      raise ValueError("Context is not set for this agent.")
-    
-    # Generating fact check or resume should not use any tools
-    message.tunables.enable_tools = False
+    async def prepare_message(self, message: Message) -> AsyncGenerator[Message, None]:
+        logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
+        if not self.context:
+            raise ValueError("Context is not set for this agent.")
 
-    async for message in super().prepare_message(message):
-      if message.status != "done":
-        yield message
+        # Generating fact check or resume should not use any tools
+        message.tunables.enable_tools = False
 
-    message.preamble["generated-resume"] = self.resume
-    job_description_agent = self.context.get_agent("job_description")
-    if not job_description_agent:
-      raise ValueError("job_description agent does not exist")
-    
-    message.preamble["job_description"] = job_description_agent.job_description
+        async for message in super().prepare_message(message):
+            if message.status != "done":
+                yield message
 
-    excluded = {}
-    preamble_types = [f"<|{p}|>" for p in message.preamble.keys() if p not in excluded]
-    preamble_types_AND = " and ".join(preamble_types)
-    preamble_types_OR = " or ".join(preamble_types)
-    message.preamble["rules"] = f"""\
+        message.preamble["generated-resume"] = self.resume
+        job_description_agent = self.context.get_agent("job_description")
+        if not job_description_agent:
+            raise ValueError("job_description agent does not exist")
+
+        message.preamble["job_description"] = job_description_agent.job_description
+
+        excluded = {}
+        preamble_types = [
+            f"<|{p}|>" for p in message.preamble.keys() if p not in excluded
+        ]
+        preamble_types_AND = " and ".join(preamble_types)
+        preamble_types_OR = " or ".join(preamble_types)
+        message.preamble[
+            "rules"
+        ] = f"""\
 - Answer the question based on the information provided in the {preamble_types_AND} sections by incorporate it seamlessly and refer to it using natural language instead of mentioning {preamble_types_OR} or quoting it directly.
 - If there is no information in these sections, answer based on your knowledge, or use any available tools.
 - Avoid phrases like 'According to the {preamble_types[0]}' or similar references to the {preamble_types_OR}.
 """
-    fact_check_agent = self.context.get_agent(agent_type="fact_check")
-    if fact_check_agent:
-      message.preamble["question"] = "Respond to:"
-    else:
-      message.preamble["question"] = f"Fact check the <|generated-resume|> based on the <|resume|>{' and <|context|>' if 'context' in message.preamble else ''}."
+        fact_check_agent = self.context.get_agent(agent_type="fact_check")
+        if fact_check_agent:
+            message.preamble["question"] = "Respond to:"
+        else:
+            message.preamble["question"] = (
+                f"Fact check the <|generated-resume|> based on the <|resume|>{' and <|context|>' if 'context' in message.preamble else ''}."
+            )
 
-    yield message
-    return
-
-  async def process_message(self, llm: Any, model: str, message:Message) -> AsyncGenerator[Message, None]:
-    logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
-    if not self.context:
-      raise ValueError("Context is not set for this agent.")
-    
-    async for message in super().process_message(llm, model, message):
-      if message.status != "done":
         yield message
+        return
 
-    fact_check_agent = self.context.get_agent(agent_type="fact_check")
-    if not fact_check_agent:
-      # Switch agent from "Fact Check Generated Resume" mode
-      # to "Answer Questions about Generated Resume"
-      self.system_prompt = system_resume
+    async def process_message(
+        self, llm: Any, model: str, message: Message
+    ) -> AsyncGenerator[Message, None]:
+        logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
+        if not self.context:
+            raise ValueError("Context is not set for this agent.")
 
-      # Instantiate the "resume" agent, and seed (or reset) its conversation
-      # with this message.
-      fact_check_agent = self.context.get_or_create_agent(agent_type="fact_check", facts=message.response)
-      first_fact_check_message = message.copy()
-      first_fact_check_message.prompt = "Fact check the generated resume."
-      fact_check_agent.conversation.add(first_fact_check_message)
-      message.response = "Resume fact checked."
+        async for message in super().process_message(llm, model, message):
+            if message.status != "done":
+                yield message
+
+        fact_check_agent = self.context.get_agent(agent_type="fact_check")
+        if not fact_check_agent:
+            # Switch agent from "Fact Check Generated Resume" mode
+            # to "Answer Questions about Generated Resume"
+            self.system_prompt = system_resume
+
+            # Instantiate the "resume" agent, and seed (or reset) its conversation
+            # with this message.
+            fact_check_agent = self.context.get_or_create_agent(
+                agent_type="fact_check", facts=message.response
+            )
+            first_fact_check_message = message.copy()
+            first_fact_check_message.prompt = "Fact check the generated resume."
+            fact_check_agent.conversation.add(first_fact_check_message)
+            message.response = "Resume fact checked."
+
+        # Return the final message
+        yield message
+        return
 
-    # Return the final message
-    yield message
-    return
 
 # Register the base agent
 agent_registry.register(Resume._agent_type, Resume)
diff --git a/src/utils/agents/types.py b/src/utils/agents/types.py
index ba9f66b..d5d87e8 100644
--- a/src/utils/agents/types.py
+++ b/src/utils/agents/types.py
@@ -1,31 +1,34 @@
 from __future__ import annotations
 from typing import List, Dict, Optional, Type
 
+
 # We'll use a registry pattern rather than hardcoded strings
 class AgentRegistry:
-  """Registry for agent types and classes"""
-  _registry: Dict[str, Type] = {}
-  
-  @classmethod
-  def register(cls, agent_type: str, agent_class: Type) -> Type:
-    """Register an agent class with its type"""
-    cls._registry[agent_type] = agent_class
-    return agent_class
-  
-  @classmethod
-  def get_class(cls, agent_type: str) -> Optional[Type]:
-    """Get the class for a given agent type"""
-    return cls._registry.get(agent_type)
-  
-  @classmethod
-  def get_types(cls) -> List[str]:
-    """Get all registered agent types"""
-    return list(cls._registry.keys())
-  
-  @classmethod
-  def get_classes(cls) -> Dict[str, Type]:
-    """Get all registered agent classes"""
-    return cls._registry.copy()
+    """Registry for agent types and classes"""
+
+    _registry: Dict[str, Type] = {}
+
+    @classmethod
+    def register(cls, agent_type: str, agent_class: Type) -> Type:
+        """Register an agent class with its type"""
+        cls._registry[agent_type] = agent_class
+        return agent_class
+
+    @classmethod
+    def get_class(cls, agent_type: str) -> Optional[Type]:
+        """Get the class for a given agent type"""
+        return cls._registry.get(agent_type)
+
+    @classmethod
+    def get_types(cls) -> List[str]:
+        """Get all registered agent types"""
+        return list(cls._registry.keys())
+
+    @classmethod
+    def get_classes(cls) -> Dict[str, Type]:
+        """Get all registered agent classes"""
+        return cls._registry.copy()
+
 
 # Create a singleton instance
-agent_registry = AgentRegistry()
\ No newline at end of file
+agent_registry = AgentRegistry()
diff --git a/src/utils/chroma.py b/src/utils/chroma.py
deleted file mode 100644
index 0fde18c..0000000
--- a/src/utils/chroma.py
+++ /dev/null
@@ -1,122 +0,0 @@
-import chromadb
-from typing import List, Dict, Any, Union
-from . import defines
-from .chunk import chunk_document
-import ollama
-
-def init_chroma_client(persist_directory: str = defines.persist_directory):
-    """Initialize and return a ChromaDB client."""
-#    return chromadb.PersistentClient(path=persist_directory)
-    return chromadb.Client()
-
-def create_or_get_collection(db: chromadb.Client, collection_name: str):
-    """Create or get a ChromaDB collection."""
-    try:
-        return db.get_collection(
-            name=collection_name
-        )
-    except:
-        return db.create_collection(
-            name=collection_name,
-            metadata={"hnsw:space": "cosine"}
-        )
-
-def process_documents_to_chroma(
-    client: ollama.Client,
-    documents: List[Dict[str, Any]],
-    collection_name: str = "document_collection",
-    text_key: str = "text",
-    max_tokens: int = 512,
-    overlap: int = 50,
-    model: str = defines.encoding_model,
-    persist_directory: str = defines.persist_directory
-):
-    """
-    Process documents, chunk them, compute embeddings, and store in ChromaDB.
-    
-    Args:
-        documents: List of document dictionaries
-        collection_name: Name for the ChromaDB collection
-        text_key: The key containing text content
-        max_tokens: Maximum tokens per chunk
-        overlap: Token overlap between chunks
-        model: Ollama model for embeddings
-        persist_directory: Directory to store ChromaDB data
-    """
-    # Initialize ChromaDB client and collection
-    db = init_chroma_client(persist_directory)
-    collection = create_or_get_collection(db, collection_name)
-    
-    # Process each document
-    for doc in documents:
-        # Chunk the document
-        doc_chunks = chunk_document(doc, text_key, max_tokens, overlap)
-        
-        # Prepare data for ChromaDB
-        ids = []
-        texts = []
-        metadatas = []
-        embeddings = []
-
-        for chunk in doc_chunks:
-            # Create a unique ID for the chunk
-            chunk_id = f"{chunk['id']}_{chunk['chunk_id']}"
-            
-            # Extract text
-            text = chunk[text_key]
-            
-            # Create metadata (excluding text and embedding to avoid duplication)
-            metadata = {k: v for k, v in chunk.items() if k != text_key and k != "embedding"}
-
-            response = client.embed(model=model, input=text)
-            embedding = response["embeddings"][0]
-            ids.append(chunk_id)
-            texts.append(text)
-            metadatas.append(metadata)
-            embeddings.append(embedding)
-        
-        # Add chunks to ChromaDB collection
-        collection.add(
-            ids=ids,
-            documents=texts,
-            embeddings=embeddings,
-            metadatas=metadatas
-        )
-    
-    return collection
-
-def query_chroma(
-    client: ollama.Client,
-    query_text: str,
-    collection_name: str = "document_collection",
-    n_results: int = 5,
-    model: str = defines.encoding_model,
-    persist_directory: str = defines.persist_directory
-):
-    """
-    Query ChromaDB for similar documents.
-    
-    Args:
-        query_text: The text to search for
-        collection_name: Name of the ChromaDB collection
-        n_results: Number of results to return
-        model: Ollama model for embedding the query
-        persist_directory: Directory where ChromaDB data is stored
-    
-    Returns:
-        Query results from ChromaDB
-    """
-    # Initialize ChromaDB client and collection
-    db = init_chroma_client(persist_directory)
-    collection = create_or_get_collection(db, collection_name)
-    
-    query_response = client.embed(model=model, input=query_text)
-    query_embeddings = query_response["embeddings"]
-
-    # Query the collection
-    results = collection.query(
-        query_embeddings=query_embeddings,
-        n_results=n_results
-    )
-    
-    return results
diff --git a/src/utils/chunk.py b/src/utils/chunk.py
deleted file mode 100644
index 636f427..0000000
--- a/src/utils/chunk.py
+++ /dev/null
@@ -1,88 +0,0 @@
-import tiktoken # type: ignore
-from . import defines
-from typing import List, Dict, Any, Union
-
-def get_encoding(model=defines.model):
-    """Get the tokenizer for counting tokens."""
-    try:
-        return tiktoken.get_encoding("cl100k_base")  # Default encoding used by many embedding models
-    except:
-        return tiktoken.encoding_for_model(model)
-
-def count_tokens(text: str) -> int:
-    """Count the number of tokens in a text string."""
-    encoding = get_encoding()
-    return len(encoding.encode(text))
-
-def chunk_text(text: str, max_tokens: int = 512, overlap: int = 50) -> List[str]:
-    """
-    Split a text into chunks based on token count with overlap between chunks.
-    
-    Args:
-        text: The text to split into chunks
-        max_tokens: Maximum number of tokens per chunk
-        overlap: Number of tokens to overlap between chunks
-    
-    Returns:
-        List of text chunks
-    """
-    if not text or max_tokens <= 0:
-        return []
-    
-    encoding = get_encoding()
-    tokens = encoding.encode(text)
-    chunks = []
-    
-    i = 0
-    while i < len(tokens):
-        # Get the current chunk of tokens
-        chunk_end = min(i + max_tokens, len(tokens))
-        chunk_tokens = tokens[i:chunk_end]
-        chunks.append(encoding.decode(chunk_tokens))
-        
-        # Move to the next position with overlap
-        if chunk_end == len(tokens):
-            break
-        i += max_tokens - overlap
-    
-    return chunks
-
-def chunk_document(document: Dict[str, Any], 
-                  text_key: str = "text",
-                  max_tokens: int = 512, 
-                  overlap: int = 50) -> List[Dict[str, Any]]:
-    """
-    Chunk a document dictionary into multiple chunks.
-    
-    Args:
-        document: Document dictionary with metadata and text
-        text_key: The key in the document that contains the text to chunk
-        max_tokens: Maximum number of tokens per chunk
-        overlap: Number of tokens to overlap between chunks
-        
-    Returns:
-        List of document dictionaries, each with chunked text and preserved metadata
-    """
-    if text_key not in document:
-        raise Exception(f"{text_key} not in document")
-    
-    # Extract text and create chunks
-    if "title" in document:
-        text = f"{document["title"]}: {document[text_key]}"
-    else:
-        text = document[text_key]
-    chunks = chunk_text(text, max_tokens, overlap)
-    
-    # Create document chunks with preserved metadata
-    chunked_docs = []
-    for i, chunk in enumerate(chunks):
-        # Create a new doc with all original fields
-        doc_chunk = document.copy()
-        # Replace text with the chunk
-        doc_chunk[text_key] = chunk
-        # Add chunk metadata
-        doc_chunk["chunk_id"] = i
-        doc_chunk["chunk_total"] = len(chunks)
-        chunked_docs.append(doc_chunk)
-    
-    return chunked_docs
diff --git a/src/utils/context.py b/src/utils/context.py
index 4db03e8..94ae19d 100644
--- a/src/utils/context.py
+++ b/src/utils/context.py
@@ -1,32 +1,35 @@
 from __future__ import annotations
-from pydantic import BaseModel, Field, model_validator# type: ignore
+from pydantic import BaseModel, Field, model_validator  # type: ignore
 from uuid import uuid4
 from typing import List, Optional, Generator, ClassVar, Any
 from typing_extensions import Annotated, Union
-import numpy as np # type: ignore
+import numpy as np  # type: ignore
 import logging
 from uuid import uuid4
-from prometheus_client import CollectorRegistry, Counter # type: ignore
+from prometheus_client import CollectorRegistry, Counter  # type: ignore
 
-from . message import Message, Tunables
-from . rag import ChromaDBFileWatcher
+from .message import Message, Tunables
+from .rag import ChromaDBFileWatcher
 from . import defines
 from . import tools as Tools
-from . agents import AnyAgent
+from .agents import AnyAgent
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+
 class Context(BaseModel):
-    model_config = {"arbitrary_types_allowed": True} # Allow ChromaDBFileWatcher
+    model_config = {"arbitrary_types_allowed": True}  # Allow ChromaDBFileWatcher
     # Required fields
     file_watcher: Optional[ChromaDBFileWatcher] = Field(default=None, exclude=True)
-    prometheus_collector: Optional[CollectorRegistry] = Field(default=None, exclude=True)
+    prometheus_collector: Optional[CollectorRegistry] = Field(
+        default=None, exclude=True
+    )
 
     # Optional fields
     id: str = Field(
         default_factory=lambda: str(uuid4()),
-        pattern=r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$"
+        pattern=r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$",
     )
     user_resume: Optional[str] = None
     user_job_description: Optional[str] = None
@@ -35,10 +38,10 @@ class Context(BaseModel):
     rags: List[dict] = []
     message_history_length: int = 5
     # Class managed fields
-    agents: List[Annotated[Union[*Agent.__subclasses__()], Field(discriminator="agent_type")]] = Field( # type: ignore
+    agents: List[Annotated[Union[*Agent.__subclasses__()], Field(discriminator="agent_type")]] = Field(  # type: ignore
         default_factory=list
     )
-    
+
     processing: bool = Field(default=False, exclude=True)
 
     # @model_validator(mode="before")
@@ -53,25 +56,29 @@ class Context(BaseModel):
         logger.info(f"Context {self.id} initialized with {len(self.agents)} agents.")
         agent_types = [agent.agent_type for agent in self.agents]
         if len(agent_types) != len(set(agent_types)):
-            raise ValueError("Context cannot contain multiple agents of the same agent_type")
+            raise ValueError(
+                "Context cannot contain multiple agents of the same agent_type"
+            )
         # for agent in self.agents:
         #   agent.set_context(self)
         return self
 
-    def generate_rag_results(self, message: Message) -> Generator[Message, None, None]:
+    def generate_rag_results(
+        self, message: Message, top_k=10, threshold=0.7
+    ) -> Generator[Message, None, None]:
         """
         Generate RAG results for the given query.
-        
+
         Args:
             query: The query string to generate RAG results for.
-        
+
         Returns:
             A list of dictionaries containing the RAG results.
         """
         try:
             message.status = "processing"
 
-            entries : int = 0
+            entries: int = 0
 
             if not self.file_watcher:
                 message.response = "No RAG context available."
@@ -86,53 +93,69 @@ class Context(BaseModel):
                     continue
                 message.response = f"Checking RAG context {rag['name']}..."
                 yield message
-                chroma_results = self.file_watcher.find_similar(query=message.prompt, top_k=10, threshold=0.7)
+                chroma_results = self.file_watcher.find_similar(
+                    query=message.prompt, top_k=top_k, threshold=threshold
+                )
                 if chroma_results:
                     entries += len(chroma_results["documents"])
 
-                    chroma_embedding = np.array(chroma_results["query_embedding"]).flatten()  # Ensure correct shape
+                    chroma_embedding = np.array(
+                        chroma_results["query_embedding"]
+                    ).flatten()  # Ensure correct shape
                     print(f"Chroma embedding shape: {chroma_embedding.shape}")
 
-                    umap_2d = self.file_watcher.umap_model_2d.transform([chroma_embedding])[0].tolist()
-                    print(f"UMAP 2D output: {umap_2d}, length: {len(umap_2d)}")  # Debug output
+                    umap_2d = self.file_watcher.umap_model_2d.transform(
+                        [chroma_embedding]
+                    )[0].tolist()
+                    print(
+                        f"UMAP 2D output: {umap_2d}, length: {len(umap_2d)}"
+                    )  # Debug output
 
-                    umap_3d = self.file_watcher.umap_model_3d.transform([chroma_embedding])[0].tolist()
-                    print(f"UMAP 3D output: {umap_3d}, length: {len(umap_3d)}")  # Debug output
+                    umap_3d = self.file_watcher.umap_model_3d.transform(
+                        [chroma_embedding]
+                    )[0].tolist()
+                    print(
+                        f"UMAP 3D output: {umap_3d}, length: {len(umap_3d)}"
+                    )  # Debug output
 
-                    message.metadata["rag"].append({
-                        "name": rag["name"],
-                        **chroma_results,
-                        "umap_embedding_2d": umap_2d,
-                        "umap_embedding_3d": umap_3d
-                    })
+                    message.metadata["rag"].append(
+                        {
+                            "name": rag["name"],
+                            **chroma_results,
+                            "umap_embedding_2d": umap_2d,
+                            "umap_embedding_3d": umap_3d,
+                        }
+                    )
                     message.response = f"Results from {rag['name']} RAG: {len(chroma_results['documents'])} results."
                     yield message
-          
+
             if entries == 0:
                 del message.metadata["rag"]
 
-            message.response = f"RAG context gathered from results from {entries} documents."
+            message.response = (
+                f"RAG context gathered from results from {entries} documents."
+            )
             message.status = "done"
             yield message
             return
         except Exception as e:
-          message.status = "error"
-          message.response = f"Error generating RAG results: {str(e)}"
-          logger.error(e)
-          yield message
-          return
-    
+            message.status = "error"
+            message.response = f"Error generating RAG results: {str(e)}"
+            logger.error(e)
+            yield message
+            return
+
     def get_or_create_agent(self, agent_type: str, **kwargs) -> Agent:
         """
         Get or create and append a new agent of the specified type, ensuring only one agent per type exists.
-        
+
         Args:
             agent_type: The type of agent to create (e.g., 'web', 'database').
             **kwargs: Additional fields required by the specific agent subclass.
-        
+
         Returns:
             The created agent instance.
-        
+
         Raises:
             ValueError: If no matching agent type is found or if a agent of this type already exists.
         """
@@ -140,7 +163,7 @@ class Context(BaseModel):
         for agent in self.agents:
             if agent.agent_type == agent_type:
                 return agent
-        
+
         # Find the matching subclass
         for agent_cls in Agent.__subclasses__():
             if agent_cls.model_fields["agent_type"].default == agent_type:
@@ -150,13 +173,15 @@ class Context(BaseModel):
                 agent.set_context(self)
                 self.agents.append(agent)
                 return agent
-        
+
         raise ValueError(f"No agent class found for agent_type: {agent_type}")
 
     def add_agent(self, agent: AnyAgent) -> None:
         """Add a Agent to the context, ensuring no duplicate agent_type."""
         if any(s.agent_type == agent.agent_type for s in self.agents):
-            raise ValueError(f"A agent with agent_type '{agent.agent_type}' already exists")
+            raise ValueError(
+                f"A agent with agent_type '{agent.agent_type}' already exists"
+            )
         self.agents.append(agent)
 
     def get_agent(self, agent_type: str) -> Agent | None:
@@ -188,5 +213,7 @@ class Context(BaseModel):
                 summary += f"\nChat Name: {agent.name}\n"
         return summary
 
-from . agents import Agent
-Context.model_rebuild()
\ No newline at end of file
+
+from .agents import Agent
+
+Context.model_rebuild()
diff --git a/src/utils/conversation.py b/src/utils/conversation.py
index 84c63f8..ff8ffe2 100644
--- a/src/utils/conversation.py
+++ b/src/utils/conversation.py
@@ -1,7 +1,8 @@
-from pydantic import BaseModel, Field, PrivateAttr # type: ignore
+from pydantic import BaseModel, Field, PrivateAttr  # type: ignore
 from typing import List
 from .message import Message
 
+
 class Conversation(BaseModel):
     Conversation_messages: List[Message] = Field(default=[], alias="messages")
 
@@ -12,24 +13,28 @@ class Conversation(BaseModel):
         return iter(self.Conversation_messages)
 
     def reset(self):
-       self.Conversation_messages = []
+        self.Conversation_messages = []
 
     @property
     def messages(self):
         """Return a copy of messages to prevent modification of the internal list."""
-        raise AttributeError("Cannot directly get messages. Use Conversation.add() or .reset()")
-    
+        raise AttributeError(
+            "Cannot directly get messages. Use Conversation.add() or .reset()"
+        )
+
     @messages.setter
     def messages(self, value):
         """Control how messages can be set, or prevent setting altogether."""
-        raise AttributeError("Cannot directly set messages. Use Conversation.add() or .reset()")
-    
+        raise AttributeError(
+            "Cannot directly set messages. Use Conversation.add() or .reset()"
+        )
+
     def add(self, message: Message | List[Message]) -> None:
         """Add a Message(s) to the conversation."""
         if isinstance(message, Message):
-          self.Conversation_messages.append(message)
+            self.Conversation_messages.append(message)
         else:
-          self.Conversation_messages.extend(message)
+            self.Conversation_messages.extend(message)
 
     def get_summary(self) -> str:
         """Return a summary of the conversation."""
@@ -38,4 +43,4 @@ class Conversation(BaseModel):
         summary = f"Conversation:\n"
         for i, message in enumerate(self.Conversation_messages, 1):
             summary += f"\nMessage {i}:\n{message.get_summary()}\n"
-        return summary
\ No newline at end of file
+        return summary
diff --git a/src/utils/defines.py b/src/utils/defines.py
index de7194e..05b2edd 100644
--- a/src/utils/defines.py
+++ b/src/utils/defines.py
@@ -1,15 +1,15 @@
 import os
 
-ollama_api_url="http://ollama:11434"  # Default Ollama local endpoint
-#model = "deepseek-r1:7b" # Tool calls don"t work
-#model="mistral:7b"       # Tool calls don"t work
-#model = "llama3.2"
-#model = "qwen3:8b"       # Requires newer ollama
-#model = "gemma3:4b"      # Requires newer ollama
+ollama_api_url = "http://ollama:11434"  # Default Ollama local endpoint
+# model = "deepseek-r1:7b" # Tool calls don"t work
+# model="mistral:7b"       # Tool calls don"t work
+# model = "llama3.2"
+# model = "qwen3:8b"       # Requires newer ollama
+# model = "gemma3:4b"      # Requires newer ollama
 model = os.getenv("MODEL_NAME", "qwen2.5:7b")
 embedding_model = os.getenv("EMBEDDING_MODEL_NAME", "mxbai-embed-large")
 persist_directory = os.getenv("PERSIST_DIR", "/opt/backstory/chromadb")
-max_context = 2048*8*2
+max_context = 2048 * 8 * 2
 doc_dir = "/opt/backstory/docs/"
 context_dir = "/opt/backstory/sessions"
 static_content = "/opt/backstory/frontend/deployed"
@@ -17,4 +17,4 @@ resume_doc = "/opt/backstory/docs/resume/resume.md"
 # Only used for testing; backstory-prod will not use this
 key_path = "/opt/backstory/keys/key.pem"
 cert_path = "/opt/backstory/keys/cert.pem"
-logging_level = os.getenv("LOGGING_LEVEL", "INFO").upper()
\ No newline at end of file
+logging_level = os.getenv("LOGGING_LEVEL", "INFO").upper()
diff --git a/src/utils/full.py b/src/utils/full.py
deleted file mode 100644
index 9affe1f..0000000
--- a/src/utils/full.py
+++ /dev/null
@@ -1,468 +0,0 @@
-import requests
-from typing import List, Dict, Any, Union
-import tiktoken
-import feedparser
-import logging as log
-import datetime
-from bs4 import BeautifulSoup
-import chromadb
-import ollama
-import re
-import numpy as np
-from . import chunk
-
-OLLAMA_API_URL = "http://ollama:11434"  # Default Ollama local endpoint
-#MODEL_NAME = "deepseek-r1:1.5b"
-MODEL_NAME = "deepseek-r1:7b"
-EMBED_MODEL = "mxbai-embed-large"
-PERSIST_DIRECTORY = "/root/.cache/chroma"
-
-client = ollama.Client(host=OLLAMA_API_URL)
-
-def extract_text_from_html_or_xml(content, is_xml=False):
-    # Parse the content
-    if is_xml:
-        soup = BeautifulSoup(content, 'xml')  # Use 'xml' parser for XML content
-    else:
-        soup = BeautifulSoup(content, 'html.parser')  # Default to 'html.parser' for HTML content
-
-    # Extract and return just the text
-    return soup.get_text()
-
-class Feed():
-    def __init__(self, name, url, poll_limit_min = 30, max_articles=5):
-        self.name = name
-        self.url = url
-        self.poll_limit_min = datetime.timedelta(minutes=poll_limit_min)
-        self.last_poll = None
-        self.articles = []
-        self.max_articles = max_articles
-        self.update()
-
-    def update(self):
-        now = datetime.datetime.now()
-        if self.last_poll is None or (now - self.last_poll) >= self.poll_limit_min:
-            log.info(f"Updating {self.name}")
-            feed = feedparser.parse(self.url)
-            self.articles = []
-            self.last_poll = now
-
-            if len(feed.entries) == 0:
-                return
-            
-            for i, entry in enumerate(feed.entries[:self.max_articles]):
-                content = {}
-                content['source'] = self.name
-                content['id'] = f"{self.name}{i}"
-                title = entry.get("title")
-                if title:
-                    content['title'] = title
-                link = entry.get("link")
-                if link:
-                    content['link'] = link
-                text = entry.get("summary")
-                if text:
-                    content['text'] = extract_text_from_html_or_xml(text, False)
-                else:
-                    continue
-                published = entry.get("published")
-                if published:
-                    content['published'] = published
-
-                self.articles.append(content)
-        else:
-            log.info(f"Not updating {self.name} -- {self.poll_limit_min - (now - self.last_poll)}s remain to refresh.")
-        return self.articles
-
-# News RSS Feeds
-rss_feeds = [
-    Feed(name="IGN.com", url="https://feeds.feedburner.com/ign/games-all"),
-    Feed(name="BBC World", url="http://feeds.bbci.co.uk/news/world/rss.xml"),
-    Feed(name="Reuters World", url="http://feeds.reuters.com/Reuters/worldNews"),
-    Feed(name="Al Jazeera", url="https://www.aljazeera.com/xml/rss/all.xml"),
-    Feed(name="CNN World", url="http://rss.cnn.com/rss/edition_world.rss"),
-    Feed(name="Time", url="https://time.com/feed/"),
-    Feed(name="Euronews", url="https://www.euronews.com/rss"),
-#    Feed(name="FeedX", url="https://feedx.net/rss/ap.xml")
-]
-
-
-def init_chroma_client(persist_directory: str = PERSIST_DIRECTORY):
-    """Initialize and return a ChromaDB client."""
-#    return chromadb.PersistentClient(path=persist_directory)
-    return chromadb.Client()
-
-def create_or_get_collection(client, collection_name: str):
-    """Create or get a ChromaDB collection."""
-    try:
-        return client.get_collection(
-            name=collection_name
-        )
-    except:
-        return client.create_collection(
-            name=collection_name,
-            metadata={"hnsw:space": "cosine"}
-        )
-
-def process_documents_to_chroma(
-    documents: List[Dict[str, Any]],
-    collection_name: str = "document_collection",
-    text_key: str = "text",
-    max_tokens: int = 512,
-    overlap: int = 50,
-    model: str = EMBED_MODEL,
-    persist_directory: str = PERSIST_DIRECTORY
-):
-    """
-    Process documents, chunk them, compute embeddings, and store in ChromaDB.
-    
-    Args:
-        documents: List of document dictionaries
-        collection_name: Name for the ChromaDB collection
-        text_key: The key containing text content
-        max_tokens: Maximum tokens per chunk
-        overlap: Token overlap between chunks
-        model: Ollama model for embeddings
-        persist_directory: Directory to store ChromaDB data
-    """
-    # Initialize ChromaDB client and collection
-    db = init_chroma_client(persist_directory)
-    collection = create_or_get_collection(db, collection_name)
-    
-    # Process each document
-    for doc in documents:
-        # Chunk the document
-        doc_chunks = chunk_document(doc, text_key, max_tokens, overlap)
-        
-        # Prepare data for ChromaDB
-        ids = []
-        texts = []
-        metadatas = []
-        embeddings = []
-
-        for chunk in doc_chunks:
-            # Create a unique ID for the chunk
-            chunk_id = f"{chunk['id']}_{chunk['chunk_id']}"
-            
-            # Extract text
-            text = chunk[text_key]
-            
-            # Create metadata (excluding text and embedding to avoid duplication)
-            metadata = {k: v for k, v in chunk.items() if k != text_key and k != "embedding"}
-
-            response = client.embed(model=model, input=text)
-            embedding = response["embeddings"][0]
-            ids.append(chunk_id)
-            texts.append(text)
-            metadatas.append(metadata)
-            embeddings.append(embedding)
-        
-        # Add chunks to ChromaDB collection
-        collection.add(
-            ids=ids,
-            documents=texts,
-            embeddings=embeddings,
-            metadatas=metadatas
-        )
-    
-    return collection
-
-def query_chroma(
-    query_text: str,
-    collection_name: str = "document_collection",
-    n_results: int = 5,
-    model: str = EMBED_MODEL,
-    persist_directory: str = PERSIST_DIRECTORY
-):
-    """
-    Query ChromaDB for similar documents.
-    
-    Args:
-        query_text: The text to search for
-        collection_name: Name of the ChromaDB collection
-        n_results: Number of results to return
-        model: Ollama model for embedding the query
-        persist_directory: Directory where ChromaDB data is stored
-    
-    Returns:
-        Query results from ChromaDB
-    """
-    # Initialize ChromaDB client and collection
-    db = init_chroma_client(persist_directory)
-    collection = create_or_get_collection(db, collection_name)
-    
-    query_response = client.embed(model=model, input=query_text)
-    query_embeddings = query_response["embeddings"]
-
-    # Query the collection
-    results = collection.query(
-        query_embeddings=query_embeddings,
-        n_results=n_results
-    )
-    
-    return results
-
-def print_top_match(query_results, index=0, documents=None):
-    """
-    Print detailed information about the top matching document,
-    including the full original document content.
-    
-    Args:
-        query_results: Results from ChromaDB query
-        documents: Original documents dictionary to look up full content (optional)
-    """
-    if not query_results or not query_results["ids"] or len(query_results["ids"][0]) == 0:
-        print("No matching documents found.")
-        return
-    
-    # Get the top result
-    top_id = query_results["ids"][0][index]
-    top_document_chunk = query_results["documents"][0][index]
-    top_metadata = query_results["metadatas"][0][index]
-    top_distance = query_results["distances"][0][index]
-    
-    print("="*50)
-    print("MATCHING DOCUMENT")
-    print("="*50)
-    print(f"Chunk ID: {top_id}")
-    print(f"Similarity Score: {top_distance:.4f}")  # Convert distance to similarity
-    
-    print("\nCHUNK METADATA:")
-    for key, value in top_metadata.items():
-        print(f"  {key}: {value}")
-    
-    print("\nMATCHING CHUNK CONTENT:")
-    print(top_document_chunk[:500].strip() + ("..." if len(top_document_chunk) > 500 else ""))
-    
-    # Extract the original document ID from the chunk ID
-    # Chunk IDs are in format "doc_id_chunk_num"
-    original_doc_id = top_id.split('_')[0]
-
-def get_top_match(query_results, index=0, documents=None):
-    top_id = query_results["ids"][index][0]
-    # Extract the original document ID from the chunk ID
-    # Chunk IDs are in format "doc_id_chunk_num"
-    original_doc_id = top_id.split('_')[0]
-        
-    # Return the full document for further processing if needed
-    if documents is not None:
-        return next((doc for doc in documents if doc["id"] == original_doc_id), None)
-    
-    return None
-
-def show_documents(documents=None):
-    if not documents:
-        return
-    
-    # Print the top matching document
-    for i, doc in enumerate(documents):
-        print(f"Document {i+1}:")
-        print(f"  Title: {doc['title']}")
-        print(f"  Text: {doc['text'][:100]}...")
-        print()
-
-def show_headlines(documents=None):
-    if not documents:
-        return
-    
-    # Print the top matching document
-    for doc in documents:
-        print(f"{doc['source']}: {doc['title']}")
-
-def show_help():
-    print("""help>
-docs       Show RAG docs
-full       Show last full top match
-headlines  Show the RAG headlines
-prompt     Show the last prompt
-response   Show the last response
-scores     Show last RAG scores
-why|think        Show last response's <think>
-context|match    Show RAG match info to last prompt
-""")
-
-
-# Example usage
-if __name__ == "__main__":
-    documents = []
-    for feed in rss_feeds:
-        documents.extend(feed.articles)
-
-    show_documents(documents=documents)
-
-    # Process documents and store in ChromaDB
-    collection = process_documents_to_chroma(
-        documents=documents,
-        collection_name="research_papers",
-        max_tokens=256,
-        overlap=25,
-        model=EMBED_MODEL,
-        persist_directory="/root/.cache/chroma"
-    )
-
-    last_results = None
-    last_prompt = None
-    last_system = None
-    last_response = None
-    last_why = None
-    last_messages = []
-    while True:
-        try:
-            search_query = input("> ").strip()
-        except KeyboardInterrupt as e:
-            print("\nExiting.")
-            break
-
-        if search_query == "exit" or search_query == "quit":
-            print("\nExiting.")
-            break
-    
-        if search_query == "docs":
-            show_documents(documents)
-            continue
-
-        if search_query == "prompt":
-            if last_prompt:
-                print(f"""last prompt>
-{"="*10}system{"="*10}
-{last_system}
-{"="*10}prompt{"="*10}
-{last_prompt}""")
-            else:
-                print(f"No prompts yet")
-            continue
-
-        if search_query == "response":
-            if last_response:
-                print(f"""last response>
-{"="*10}response{"="*10}
-{last_response}""")
-            else:
-                print(f"No responses yet")
-            continue
-
-        if search_query == "" or search_query == "help":
-            show_help()
-            continue
-
-        if search_query == "headlines":
-            show_headlines(documents)
-            continue
-
-        if search_query == "match" or search_query == "context":
-            if last_results:
-                print_top_match(last_results, documents=documents)
-            else:
-                print("No match to give info on")
-            continue
-
-        if search_query == "why" or search_query == "think":
-            if last_why:
-                print(f"""
-why>
-{last_why}
-""")
-            else:
-                print("No processed prompts")
-            continue
-
-        if search_query == "scores":
-            if last_results:
-                for i, _ in enumerate(last_results):
-                    print_top_match(last_results, documents=documents, index=i)
-            else:
-                print("No match to give info on")
-            continue
-
-        if search_query == "full":
-            if last_results:
-                full = get_top_match(last_results, documents=documents)
-                if full:
-                    print(f"""Context:
-Source: {full["source"]}
-Title: {full["title"]}
-Link: {full["link"]}
-Distance: {last_results.get("distances", [[0]])[0][0]}
-Full text:
-{full["text"]}""")
-            else:
-                print("No match to give info on")
-            continue
-
-        # Query ChromaDB
-        results = query_chroma(
-            query_text=search_query,
-            collection_name="research_papers",
-            n_results=10
-        )
-        last_results = results
-
-        full = get_top_match(results, documents=documents)
-
-        headlines = ""
-        for doc in documents:
-            headlines += f"{doc['source']}: {doc['title']}\n"
-
-        system=f"""
-You are the assistant. Your name is airc. This application is called airc (pronounced Eric).
-
-Information about the author of this program and the AI model it uses:
-
-* James wrote the python application called airc that is driving this RAG model on top of {MODEL_NAME} using {EMBED_MODEL} and chromadb for vector embedding. Link https://github.com/jketreno/airc.
-* James Ketrenos is a software engineer with a history in all levels of the computer stack, from the kernel to full-stack web applications. He dabbles in AI/ML and is familiar with pytorch and ollama.
-* James Ketrenos deployed this application locally on an Intel Arc B580 (battlemage) computer using Intel's ipex-llm. 
-* For Intel GPU metrics, James Ketrenos wrote the "ze-monitor" utility in C++. ze-monitor provides Intel GPU telemetry data for Intel client GPU devices, similar to xpu-smi. Link https://github.com/jketreno/ze-monitor. airc uses ze-monitor.
-* James lives in Portland, Oregon and has three kids. Two are attending Oregon State University and one is attending Williamette University.
-* airc provides an IRC chat bot as well as a React web frontend available at https://airc.ketrenos.com
-
-You must follow these rules:
-
-* Provide short (less than 100 character) responses.
-* Provide a single response.
-* Do not prefix it with a word like 'Answer'.
-* For information about the AI running this system, include information about author, including links.
-* For information relevant to the current events in the <input></input> tags, use that information and state the source when information comes from.
-
-"""
-        context = "Information related to current events\n<input>=["
-        for doc in documents:
-            item = {'source':doc["source"],'article':{'title':doc["title"],'link':doc["link"],'text':doc["text"]}}
-            context += f"{item}"
-        context += "\n</input>"
-
-        prompt = f"{search_query}"
-        last_prompt = prompt
-        last_system = system # cache it before news context is added
-        system = f"{system}{context}"
-        if len(last_messages) != 0:
-            message_context = f"{last_messages}"
-            prompt = f"{message_context}{prompt}"
-
-        print(f"system len: {len(system)}")            
-        print(f"prompt len: {len(prompt)}")            
-        output = client.generate(
-            model=MODEL_NAME,
-            system=system,
-            prompt=prompt,
-            stream=False,
-            options={ 'num_ctx': 100000 }
-        )
-        # Prune off the <think>...</think>
-        matches = re.match(r'^<think>(.*?)</think>(.*)$', output['response'], flags=re.DOTALL)
-        if matches:
-            last_why = matches[1].strip()
-            content = matches[2].strip()
-        else:
-            print(f"[garbled] response>\n{output['response']}")
-        print(f"Response>\n{content}")
-
-        last_response = content
-        last_messages.extend(({
-            'role': 'user',
-            'name': 'james',
-            'message': search_query
-        }, {
-            'role': 'assistant',
-            'message': content
-        }))
-        last_messages = last_messages[:10]
diff --git a/src/utils/message.py b/src/utils/message.py
index 21208e2..5b0a80f 100644
--- a/src/utils/message.py
+++ b/src/utils/message.py
@@ -3,55 +3,69 @@ from typing import Dict, List, Optional, Any
 from datetime import datetime, timezone
 from asyncio import Event
 
+
 class Tunables(BaseModel):
-    enable_rag : bool = Field(default=True)     # Enable RAG collection chromadb matching
-    enable_tools : bool = Field(default=True)   # Enable LLM to use tools
-    enable_context : bool = Field(default=True) # Add <|context|> field to message
+    enable_rag: bool = Field(default=True)  # Enable RAG collection chromadb matching
+    enable_tools: bool = Field(default=True)  # Enable LLM to use tools
+    enable_context: bool = Field(default=True)  # Add <|context|> field to message
+
 
 class Message(BaseModel):
-    model_config = {"arbitrary_types_allowed": True} # Allow Event
+    model_config = {"arbitrary_types_allowed": True}  # Allow Event
     # Required
-    prompt: str                     # Query to be answered
+    prompt: str  # Query to be answered
 
     # Tunables
     tunables: Tunables = Field(default_factory=Tunables)
-    
+
     # Generated while processing message
-    status: str = ""                # Status of the message
-    preamble: dict[str,str] = {}    # Preamble to be prepended to the prompt
-    system_prompt: str = ""         # System prompt provided to the LLM
-    context_prompt: str = ""        # Full content of the message (preamble + prompt)
-    response: str = ""              # LLM response to the preamble + query
-    metadata: Dict[str, Any] = Field(default_factory=lambda: {
-        "rag": [],
-        "eval_count": 0,
-        "eval_duration": 0,
-        "prompt_eval_count": 0,
-        "prompt_eval_duration": 0,
-        "context_size": 0,
-    })
-    network_packets: int = 0        # Total number of streaming packets
-    network_bytes: int = 0          # Total bytes sent while streaming packets
-    actions: List[str] = []         # Other session modifying actions performed while processing the message 
+    status: str = ""  # Status of the message
+    preamble: dict[str, str] = {}  # Preamble to be prepended to the prompt
+    system_prompt: str = ""  # System prompt provided to the LLM
+    context_prompt: str = ""  # Full content of the message (preamble + prompt)
+    response: str = ""  # LLM response to the preamble + query
+    metadata: Dict[str, Any] = Field(
+        default_factory=lambda: {
+            "rag": [],
+            "eval_count": 0,
+            "eval_duration": 0,
+            "prompt_eval_count": 0,
+            "prompt_eval_duration": 0,
+            "context_size": 0,
+        }
+    )
+    network_packets: int = 0  # Total number of streaming packets
+    network_bytes: int = 0  # Total bytes sent while streaming packets
+    actions: List[str] = (
+        []
+    )  # Other session modifying actions performed while processing the message
     timestamp: datetime = datetime.now(timezone.utc)
-    chunk: str = Field(default="")  # This needs to be serialized so it will be sent in responses
-    partial_response: str = Field(default="") # This needs to be serialized so it will be sent in responses on timeout
-    
+    chunk: str = Field(
+        default=""
+    )  # This needs to be serialized so it will be sent in responses
+    partial_response: str = Field(
+        default=""
+    )  # This needs to be serialized so it will be sent in responses on timeout
+    title: str = Field(
+        default=""
+    )  # This needs to be serialized so it will be sent in responses on timeout
+
     def add_action(self, action: str | list[str]) -> None:
         """Add a actions(s) to the message."""
         if isinstance(action, str):
-          self.actions.append(action)
+            self.actions.append(action)
         else:
-          self.actions.extend(action)
-        
+            self.actions.extend(action)
+
     def get_summary(self) -> str:
         """Return a summary of the message."""
         response_summary = (
             f"Response: {self.response} (Actions: {', '.join(self.actions)})"
-            if self.response else "No response yet"
+            if self.response
+            else "No response yet"
         )
         return (
             f"Message at {self.timestamp}:\n"
             f"Query: {self.preamble}{self.content}\n"
             f"{response_summary}"
-        )
\ No newline at end of file
+        )
diff --git a/src/utils/metrics.py b/src/utils/metrics.py
index 69d813d..3d3b111 100644
--- a/src/utils/metrics.py
+++ b/src/utils/metrics.py
@@ -1,6 +1,7 @@
-from prometheus_client import Counter, Gauge, Summary, Histogram, Info, Enum, CollectorRegistry # type: ignore
+from prometheus_client import Counter, Gauge, Summary, Histogram, Info, Enum, CollectorRegistry  # type: ignore
 from threading import Lock
 
+
 def singleton(cls):
     instance = None
     lock = Lock()
@@ -14,80 +15,81 @@ def singleton(cls):
 
     return get_instance
 
+
 @singleton
-class Metrics():
-  def __init__(self, *args, prometheus_collector, **kwargs):
-    super().__init__(*args, **kwargs)
-    self.prometheus_collector = prometheus_collector
-    
-    self.prepare_count : Counter = Counter(
-      name="prepare_total",
-      documentation="Total messages prepared by agent type",
-      labelnames=("agent",),
-      registry=self.prometheus_collector
-    )
-    
-    self.prepare_duration : Histogram = Histogram(
-      name="prepare_duration",
-      documentation="Preparation duration by agent type",
-      labelnames=("agent",),
-      registry=self.prometheus_collector
-    )
-    
-    self.process_count : Counter = Counter(
-      name="process",
-      documentation="Total messages processed by agent type",
-      labelnames=("agent",),
-      registry=self.prometheus_collector
-    )
-    
-    self.process_duration : Histogram = Histogram(
-      name="process_duration",
-      documentation="Processing duration by agent type",
-      labelnames=("agent",),
-      registry=self.prometheus_collector
-    )
-        
-    self.tool_count : Counter = Counter(
-      name="tool_total",
-      documentation="Total messages tooled by agent type",
-      labelnames=("agent",),
-      registry=self.prometheus_collector
-    )
-    
-    self.tool_duration : Histogram = Histogram(
-      name="tool_duration",
-      documentation="Tool duration by agent type",
-      buckets=(0.1, 0.5, 1.0, 2.0, float('inf')),
-      labelnames=("agent",),
-      registry=self.prometheus_collector
-    )
-    
-    self.generate_count : Counter = Counter(
-      name="generate_total",
-      documentation="Total messages generated by agent type",
-      labelnames=("agent",),
-      registry=self.prometheus_collector
-    )
-    
-    self.generate_duration : Histogram = Histogram(
-      name="generate_duration",
-      documentation="Generate duration by agent type",
-      buckets=(0.1, 0.5, 1.0, 2.0, float('inf')),
-      labelnames=("agent",),
-      registry=self.prometheus_collector
-    )
-    
-    self.tokens_prompt : Counter = Counter(
-      name="tokens_prompt",
-      documentation="Total tokens passed as prompt to LLM",
-      labelnames=("agent",),
-      registry=self.prometheus_collector
-    )
-    
-    self.tokens_eval : Counter = Counter(
-      name="tokens_eval",
-      documentation="Total tokens returned by LLM",
-      labelnames=("agent",),
-      registry=self.prometheus_collector
-    )
\ No newline at end of file
+class Metrics:
+    def __init__(self, *args, prometheus_collector, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.prometheus_collector = prometheus_collector
+
+        self.prepare_count: Counter = Counter(
+            name="prepare_total",
+            documentation="Total messages prepared by agent type",
+            labelnames=("agent",),
+            registry=self.prometheus_collector,
+        )
+
+        self.prepare_duration: Histogram = Histogram(
+            name="prepare_duration",
+            documentation="Preparation duration by agent type",
+            labelnames=("agent",),
+            registry=self.prometheus_collector,
+        )
+
+        self.process_count: Counter = Counter(
+            name="process",
+            documentation="Total messages processed by agent type",
+            labelnames=("agent",),
+            registry=self.prometheus_collector,
+        )
+
+        self.process_duration: Histogram = Histogram(
+            name="process_duration",
+            documentation="Processing duration by agent type",
+            labelnames=("agent",),
+            registry=self.prometheus_collector,
+        )
+
+        self.tool_count: Counter = Counter(
+            name="tool_total",
+            documentation="Total messages tooled by agent type",
+            labelnames=("agent",),
+            registry=self.prometheus_collector,
+        )
+
+        self.tool_duration: Histogram = Histogram(
+            name="tool_duration",
+            documentation="Tool duration by agent type",
+            buckets=(0.1, 0.5, 1.0, 2.0, float("inf")),
+            labelnames=("agent",),
+            registry=self.prometheus_collector,
+        )
+
+        self.generate_count: Counter = Counter(
+            name="generate_total",
+            documentation="Total messages generated by agent type",
+            labelnames=("agent",),
+            registry=self.prometheus_collector,
+        )
+
+        self.generate_duration: Histogram = Histogram(
+            name="generate_duration",
+            documentation="Generate duration by agent type",
+            buckets=(0.1, 0.5, 1.0, 2.0, float("inf")),
+            labelnames=("agent",),
+            registry=self.prometheus_collector,
+        )
+
+        self.tokens_prompt: Counter = Counter(
+            name="tokens_prompt",
+            documentation="Total tokens passed as prompt to LLM",
+            labelnames=("agent",),
+            registry=self.prometheus_collector,
+        )
+
+        self.tokens_eval: Counter = Counter(
+            name="tokens_eval",
+            documentation="Total tokens returned by LLM",
+            labelnames=("agent",),
+            registry=self.prometheus_collector,
+        )
diff --git a/src/utils/rag.py b/src/utils/rag.py
index d9b335b..9e26737 100644
--- a/src/utils/rag.py
+++ b/src/utils/rag.py
@@ -1,4 +1,4 @@
-from pydantic import BaseModel # type: ignore
+from pydantic import BaseModel  # type: ignore
 from typing import List, Optional, Dict, Any
 import os
 import glob
@@ -13,18 +13,22 @@ import time
 import hashlib
 import asyncio
 import json
-import numpy as np # type: ignore
+import numpy as np  # type: ignore
+import traceback
+import os
 
 import chromadb
 import ollama
-from langchain.text_splitter import CharacterTextSplitter # type: ignore
-from sentence_transformers import SentenceTransformer # type: ignore
-from langchain.schema import Document # type: ignore
-from watchdog.observers import Observer # type: ignore
-from watchdog.events import FileSystemEventHandler # type: ignore
-import umap # type: ignore
-from markitdown import MarkItDown # type: ignore
-from chromadb.api.models.Collection import Collection # type: ignore
+from watchdog.observers import Observer  # type: ignore
+from watchdog.events import FileSystemEventHandler  # type: ignore
+import umap  # type: ignore
+from markitdown import MarkItDown  # type: ignore
+from chromadb.api.models.Collection import Collection  # type: ignore
+
+from .markdown_chunker import (
+    MarkdownChunker,
+    Chunk,
+)
 
 # Import your existing modules
 if __name__ == "__main__":
@@ -34,13 +38,11 @@ else:
     # When imported as a module, use relative imports
     from . import defines
 
-__all__ = [
-    'ChromaDBFileWatcher',
-    'start_file_watcher'
-]
+__all__ = ["ChromaDBFileWatcher", "start_file_watcher"]
+
+DEFAULT_CHUNK_SIZE = 750
+DEFAULT_CHUNK_OVERLAP = 100
 
-DEFAULT_CHUNK_SIZE=750
-DEFAULT_CHUNK_OVERLAP=100
 
 class ChromaDBGetResponse(BaseModel):
     ids: List[str]
@@ -48,9 +50,19 @@ class ChromaDBGetResponse(BaseModel):
     documents: Optional[List[str]] = None
     metadatas: Optional[List[Dict[str, Any]]] = None
 
+
 class ChromaDBFileWatcher(FileSystemEventHandler):
-    def __init__(self, llm, watch_directory, loop, persist_directory=None, collection_name="documents", 
-                 chunk_size=DEFAULT_CHUNK_SIZE, chunk_overlap=DEFAULT_CHUNK_OVERLAP, recreate=False):
+    def __init__(
+        self,
+        llm,
+        watch_directory,
+        loop,
+        persist_directory=None,
+        collection_name="documents",
+        chunk_size=DEFAULT_CHUNK_SIZE,
+        chunk_overlap=DEFAULT_CHUNK_OVERLAP,
+        recreate=False,
+    ):
         self.llm = llm
         self.watch_directory = watch_directory
         self.persist_directory = persist_directory or defines.persist_directory
@@ -58,38 +70,34 @@ class ChromaDBFileWatcher(FileSystemEventHandler):
         self.chunk_size = chunk_size
         self.chunk_overlap = chunk_overlap
         self.loop = loop
-        self._umap_collection : ChromaDBGetResponse | None = None
-        self._umap_embedding_2d : np.ndarray = []
-        self._umap_embedding_3d : np.ndarray = []
-        self._umap_model_2d : umap.UMAP = None
-        self._umap_model_3d : umap.UMAP = None
-        self.md = MarkItDown(enable_plugins=False) # Set to True to enable plugins
+        self._umap_collection: ChromaDBGetResponse | None = None
+        self._umap_embedding_2d: np.ndarray = []
+        self._umap_embedding_3d: np.ndarray = []
+        self._umap_model_2d: umap.UMAP = None
+        self._umap_model_3d: umap.UMAP = None
+        self.md = MarkItDown(enable_plugins=False)  # Set to True to enable plugins
 
-        #self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
+        # self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
 
         # Path for storing file hash state
-        self.hash_state_path = os.path.join(self.persist_directory, f"{collection_name}_hash_state.json")
-        
+        self.hash_state_path = os.path.join(
+            self.persist_directory, f"{collection_name}_hash_state.json"
+        )
+
         # Flag to track if this is a new collection
         self.is_new_collection = False
-        
+
         # Initialize ChromaDB collection
-        self._collection : Collection = self._get_vector_collection(recreate=recreate)
+        self._collection: Collection = self._get_vector_collection(recreate=recreate)
+        self._markdown_chunker = MarkdownChunker()
         self._update_umaps()
 
         # Setup text splitter
-        self.text_splitter = CharacterTextSplitter(
-            chunk_size=chunk_size, 
-            chunk_overlap=chunk_overlap,
-            separator="\n\n",  # Respect paragraph/section breaks
-            length_function=len
-        )
-        
         # Track file hashes and processing state
         self.file_hashes = self._load_hash_state()
         self.update_lock = asyncio.Lock()
         self.processing_files = set()
-    
+
     @property
     def collection(self):
         return self._collection
@@ -101,11 +109,11 @@ class ChromaDBFileWatcher(FileSystemEventHandler):
     @property
     def umap_embedding_2d(self) -> np.ndarray:
         return self._umap_embedding_2d
-    
+
     @property
     def umap_embedding_3d(self) -> np.ndarray:
         return self._umap_embedding_3d
-    
+
     @property
     def umap_model_2d(self):
         return self._umap_model_2d
@@ -114,8 +122,8 @@ class ChromaDBFileWatcher(FileSystemEventHandler):
     def umap_model_3d(self):
         return self._umap_model_3d
 
-    def _markitdown(self, document : str, markdown : Path):
-        logging.info(f'Converting {document} to {markdown}')
+    def _markitdown(self, document: str, markdown: Path):
+        logging.info(f"Converting {document} to {markdown}")
         try:
             result = self.md.convert(document)
             markdown.write_text(result.text_content)
@@ -127,91 +135,105 @@ class ChromaDBFileWatcher(FileSystemEventHandler):
         try:
             # Create directory if it doesn't exist
             os.makedirs(os.path.dirname(self.hash_state_path), exist_ok=True)
-            
-            with open(self.hash_state_path, 'w') as f:
+
+            with open(self.hash_state_path, "w") as f:
                 json.dump(self.file_hashes, f)
-                
+
             logging.info(f"Saved hash state with {len(self.file_hashes)} entries")
         except Exception as e:
             logging.error(f"Error saving hash state: {e}")
-    
+
     def _load_hash_state(self):
         """Load the file hash state from disk."""
         if os.path.exists(self.hash_state_path):
             try:
-                with open(self.hash_state_path, 'r') as f:
+                with open(self.hash_state_path, "r") as f:
                     hash_state = json.load(f)
                 logging.info(f"Loaded hash state with {len(hash_state)} entries")
                 return hash_state
             except Exception as e:
                 logging.error(f"Error loading hash state: {e}")
-        
+
         return {}
-    
+
     async def scan_directory(self, process_all=False):
         """
         Scan directory for new, modified, or deleted files and update collection.
-        
+
         Args:
             process_all: If True, process all files regardless of hash status
         """
         # Check for new or modified files
-        file_paths = glob.glob(os.path.join(self.watch_directory, "**/*"), recursive=True)
+        file_paths = glob.glob(
+            os.path.join(self.watch_directory, "**/*"), recursive=True
+        )
         files_checked = 0
         files_processed = 0
         files_to_process = []
-        
+
         logging.info(f"Starting directory scan. Found {len(file_paths)} total paths.")
-        
+
         for file_path in file_paths:
             if os.path.isfile(file_path):
                 # Do not put the Resume in RAG as it is provideded with all queries.
-                if file_path == defines.resume_doc:
-                    logging.info(f"Not adding {file_path} to RAG -- primary resume")
-                    continue
+                # if file_path == defines.resume_doc:
+                #     logging.info(f"Not adding {file_path} to RAG -- primary resume")
+                #     continue
                 files_checked += 1
                 current_hash = self._get_file_hash(file_path)
                 if not current_hash:
                     continue
 
                 # If file is new, changed, or we're processing all files
-                if process_all or file_path not in self.file_hashes or self.file_hashes[file_path] != current_hash:
+                if (
+                    process_all
+                    or file_path not in self.file_hashes
+                    or self.file_hashes[file_path] != current_hash
+                ):
                     self.file_hashes[file_path] = current_hash
                     files_to_process.append(file_path)
-                    logging.info(f"File {'found' if process_all else 'changed'}: {file_path}")
-        
-        logging.info(f"Found {len(files_to_process)} files to process after scanning {files_checked} files")
-        
+                    logging.info(
+                        f"File {'found' if process_all else 'changed'}: {file_path}"
+                    )
+
+        logging.info(
+            f"Found {len(files_to_process)} files to process after scanning {files_checked} files"
+        )
+
         # Check for deleted files
         deleted_files = []
         for file_path in self.file_hashes:
             if not os.path.exists(file_path):
                 deleted_files.append(file_path)
                 # Schedule removal
-                asyncio.run_coroutine_threadsafe(self.remove_file_from_collection(file_path), self.loop)
+                asyncio.run_coroutine_threadsafe(
+                    self.remove_file_from_collection(file_path), self.loop
+                )
                 # Don't block on result, just let it run
                 logging.info(f"File deleted: {file_path}")
-        
+
         # Remove deleted files from hash state
         for file_path in deleted_files:
             del self.file_hashes[file_path]
-        
+
         # Process all discovered files using asyncio.gather with the existing loop
         if files_to_process:
             logging.info(f"Starting to process {len(files_to_process)} files")
-            
+
             for file_path in files_to_process:
                 async with self.update_lock:
-                    await self._update_document_in_collection(file_path) 
+                    await self._update_document_in_collection(file_path)
         else:
             logging.info("No files to process")
-        
+
         # Save the updated state
         self._save_hash_state()
-        
-        logging.info(f"Scan complete: Checked {files_checked} files, processed {files_processed}, removed {len(deleted_files)}")
-        return files_processed    
-    
+
+        logging.info(
+            f"Scan complete: Checked {files_checked} files, processed {files_processed}, removed {len(deleted_files)}"
+        )
+        return files_processed
+
     async def process_file_update(self, file_path):
         """Process a file update event."""
         # Skip if already being processed
@@ -219,34 +241,37 @@ class ChromaDBFileWatcher(FileSystemEventHandler):
             logging.info(f"{file_path} already in queue. Not adding.")
             return
 
-        if file_path == defines.resume_doc:
-            logging.info(f"Not adding {file_path} to RAG -- primary resume")
-            return
-            
+        # if file_path == defines.resume_doc:
+        #     logging.info(f"Not adding {file_path} to RAG -- primary resume")
+        #     return
+
         try:
             logging.info(f"{file_path} not in queue. Adding.")
             self.processing_files.add(file_path)
-            
+
             # Wait a moment to ensure the file write is complete
             await asyncio.sleep(0.5)
-            
+
             # Check if content changed via hash
             current_hash = self._get_file_hash(file_path)
             if not current_hash:  # File might have been deleted or is inaccessible
                 return
-                
-            if file_path in self.file_hashes and self.file_hashes[file_path] == current_hash:
+
+            if (
+                file_path in self.file_hashes
+                and self.file_hashes[file_path] == current_hash
+            ):
                 # File hasn't actually changed in content
                 logging.info(f"Hash has not changed for {file_path}")
                 return
-            
+
             # Update file hash
             self.file_hashes[file_path] = current_hash
 
             # Process and update the file in ChromaDB
             async with self.update_lock:
                 await self._update_document_in_collection(file_path)
-                
+
             # Save the hash state after successful update
             self._save_hash_state()
 
@@ -257,150 +282,162 @@ class ChromaDBFileWatcher(FileSystemEventHandler):
             logging.error(f"Error processing update for {file_path}: {e}")
         finally:
             self.processing_files.discard(file_path)
-    
+
     async def remove_file_from_collection(self, file_path):
         """Remove all chunks related to a deleted file."""
         async with self.update_lock:
             try:
                 # Find all documents with the specified path
-                results = self.collection.get(
-                    where={"path": file_path}
-                )
-                
-                if results and 'ids' in results and results['ids']:
-                    self.collection.delete(ids=results['ids'])
-                    logging.info(f"Removed {len(results['ids'])} chunks for deleted file: {file_path}")
-                
+                results = self.collection.get(where={"path": file_path})
+
+                if results and "ids" in results and results["ids"]:
+                    self.collection.delete(ids=results["ids"])
+                    logging.info(
+                        f"Removed {len(results['ids'])} chunks for deleted file: {file_path}"
+                    )
+
                 # Remove from hash dictionary
                 if file_path in self.file_hashes:
                     del self.file_hashes[file_path]
                     # Save the updated hash state
                     self._save_hash_state()
-                    
+
             except Exception as e:
                 logging.error(f"Error removing file from collection: {e}")
 
     def _update_umaps(self):
         # Update the UMAP embeddings
-        self._umap_collection = self._collection.get(include=["embeddings", "documents", "metadatas"])
+        self._umap_collection = self._collection.get(
+            include=["embeddings", "documents", "metadatas"]
+        )
         if not self._umap_collection or not len(self._umap_collection["embeddings"]):
             logging.warning("No embeddings found in the collection.")
             return
-        
-        # During initialization
-        logging.info(f"Updating 2D UMAP for {len(self._umap_collection['embeddings'])} vectors")
-        vectors = np.array(self._umap_collection["embeddings"])
-        self._umap_model_2d = umap.UMAP(n_components=2, random_state=8911, metric="cosine", n_neighbors=15, min_dist=0.1)
-        self._umap_embedding_2d = self._umap_model_2d.fit_transform(vectors)
-        logging.info(f"2D UMAP model n_components: {self._umap_model_2d.n_components}")  # Should be 2
 
-        logging.info(f"Updating 3D UMAP for {len(self._umap_collection['embeddings'])} vectors")
-        self._umap_model_3d = umap.UMAP(n_components=3, random_state=8911, metric="cosine", n_neighbors=15, min_dist=0.1)
+        # During initialization
+        logging.info(
+            f"Updating 2D UMAP for {len(self._umap_collection['embeddings'])} vectors"
+        )
+        vectors = np.array(self._umap_collection["embeddings"])
+        self._umap_model_2d = umap.UMAP(
+            n_components=2,
+            random_state=8911,
+            metric="cosine",
+            n_neighbors=15,
+            min_dist=0.1,
+        )
+        self._umap_embedding_2d = self._umap_model_2d.fit_transform(vectors)
+        logging.info(
+            f"2D UMAP model n_components: {self._umap_model_2d.n_components}"
+        )  # Should be 2
+
+        logging.info(
+            f"Updating 3D UMAP for {len(self._umap_collection['embeddings'])} vectors"
+        )
+        self._umap_model_3d = umap.UMAP(
+            n_components=3,
+            random_state=8911,
+            metric="cosine",
+            n_neighbors=15,
+            min_dist=0.1,
+        )
         self._umap_embedding_3d = self._umap_model_3d.fit_transform(vectors)
-        logging.info(f"3D UMAP model n_components: {self._umap_model_3d.n_components}")  # Should be 3
+        logging.info(
+            f"3D UMAP model n_components: {self._umap_model_3d.n_components}"
+        )  # Should be 3
 
     def _get_vector_collection(self, recreate=False) -> Collection:
         """Get or create a ChromaDB collection."""
         # Initialize ChromaDB client
-        chroma_client = chromadb.PersistentClient( # type: ignore
-            path=self.persist_directory, 
-            settings=chromadb.Settings(anonymized_telemetry=False) # type: ignore
+        chroma_client = chromadb.PersistentClient(  # type: ignore
+            path=self.persist_directory,
+            settings=chromadb.Settings(anonymized_telemetry=False),  # type: ignore
         )
 
-        # Check if the collection exists 
+        # Check if the collection exists
         try:
             chroma_client.get_collection(self.collection_name)
             collection_exists = True
         except:
             collection_exists = False
-            
+
         # If collection doesn't exist, mark it as new
         if not collection_exists:
             self.is_new_collection = True
             logging.info(f"Creating new collection: {self.collection_name}")
-            
+
         # Delete if recreate is True
         if recreate and collection_exists:
             chroma_client.delete_collection(name=self.collection_name)
             self.is_new_collection = True
             logging.info(f"Recreating collection: {self.collection_name}")
-                
+
         return chroma_client.get_or_create_collection(
-            name=self.collection_name,
-            metadata={
-                "hnsw:space": "cosine"
-            })
-    
-    def load_text_files(self, directory=None, encoding="utf-8"):
-        """Load all text files from a directory into Document objects."""
-        directory = directory or self.watch_directory
-        file_paths = glob.glob(os.path.join(directory, "**/*"), recursive=True)
-        documents = []
+            name=self.collection_name, metadata={"hnsw:space": "cosine"}
+        )
 
-        for file_path in file_paths:
-            if os.path.isfile(file_path):  # Ensure it's a file, not a directory
-                try:
-                    with open(file_path, "r", encoding=encoding) as f:
-                        content = f.read()
-                        
-                        # Extract top-level directory
-                        rel_path = os.path.relpath(file_path, directory)
-                        top_level_dir = rel_path.split(os.sep)[0]
-
-                        documents.append(Document(
-                            page_content=content,
-                            metadata={"doc_type": top_level_dir, "path": file_path}
-                        ))
-                except Exception as e:
-                    logging.error(f"Failed to load {file_path}: {e}")
-
-        return documents
-    
     def create_chunks_from_documents(self, docs):
         """Split documents into chunks using the text splitter."""
         return self.text_splitter.split_documents(docs)
-    
+
     def get_embedding(self, text, normalize=True):
         """Generate embeddings using Ollama."""
-        #response = self.embedding_model.encode(text)  # Outputs 384-dim vectors
+        # response = self.embedding_model.encode(text)  # Outputs 384-dim vectors
+
+        response = self.llm.embeddings(model=defines.embedding_model, prompt=text)
+        embedding = response["embedding"]
 
-        response = self.llm.embeddings(
-            model=defines.embedding_model, 
-            prompt=text)
-        embedding = response['embedding']
-        
         # response = self.llm.embeddings.create(
-        #     model=defines.embedding_model, 
-        #     input=text, 
+        #     model=defines.embedding_model,
+        #     input=text,
         #     options={"num_ctx": self.chunk_size * 3} # No need waste ctx space
         # )
         if normalize:
             normalized = self._normalize_embeddings(embedding)
             return normalized
         return embedding
-    
-    def add_embeddings_to_collection(self, chunks):
+
+    def add_embeddings_to_collection(self, chunks: List[Chunk]):
         """Add embeddings for chunks to the collection."""
+
         for i, chunk in enumerate(chunks):
-            text = chunk.page_content
-            metadata = chunk.metadata
+            text = chunk["text"]
+            metadata = chunk["metadata"]
 
             # Generate a more unique ID based on content and metadata
             content_hash = hashlib.md5(text.encode()).hexdigest()
             path_hash = ""
             if "path" in metadata:
-                path_hash = hashlib.md5(metadata["path"].encode()).hexdigest()[:8]
-            
+                path_hash = hashlib.md5(metadata["source_file"].encode()).hexdigest()[
+                    :8
+                ]
+
             chunk_id = f"{path_hash}_{content_hash}_{i}"
-            
+
             embedding = self.get_embedding(text)
-            self.collection.add(
-                ids=[chunk_id], 
-                documents=[text], 
-                embeddings=[embedding],
-                metadatas=[metadata]
-            )
-    
+            try:
+                self.collection.add(
+                    ids=[chunk_id],
+                    documents=[text],
+                    embeddings=[embedding],
+                    metadatas=[metadata],
+                )
+            except Exception as e:
+                logging.error(f"Error adding chunk to collection: {e}")
+                logging.error(traceback.format_exc())
+                logging.error(chunk)
+
+    def read_line_range(self, file_path, start, end, buffer=5) -> list[str]:
+        try:
+            with open(file_path, "r") as file:
+                lines = file.readlines()
+                start = max(0, start - buffer)
+                end = min(len(lines), end + buffer)
+                return lines[start:end]
+        except:
+            logging.warning(f"Unable to open {file_path}")
+            return []
+
     # Cosine Distance  Equivalent Similarity  Retrieval Characteristics
     # 0.2 - 0.3        0.85 - 0.90            Very strict, highly precise results only
     # 0.3 - 0.5        0.75 - 0.85            Strong relevance, good precision
@@ -413,16 +450,16 @@ class ChromaDBFileWatcher(FileSystemEventHandler):
         # collection is configured with hnsw:space cosine
         query_embedding = self.get_embedding(query)
         results = self.collection.query(
-            query_embeddings=[query_embedding], 
+            query_embeddings=[query_embedding],
             n_results=top_k,
             include=["documents", "metadatas", "distances"],
         )
 
         # Extract results
-        ids = results['ids'][0]
-        documents = results['documents'][0]
-        distances = results['distances'][0]
-        metadatas = results['metadatas'][0] 
+        ids = results["ids"][0]
+        documents = results["documents"][0]
+        distances = results["distances"][0]
+        metadatas = results["metadatas"][0]
 
         filtered_ids = []
         filtered_documents = []
@@ -436,6 +473,14 @@ class ChromaDBFileWatcher(FileSystemEventHandler):
                 filtered_metadatas.append(metadatas[i])
                 filtered_distances.append(distance)
 
+        for index, meta in enumerate(filtered_metadatas):
+            source_file = meta["source_file"]
+            del meta["source_file"]
+            lines = self.read_line_range(
+                source_file, meta["line_begin"], meta["line_end"]
+            )
+            if len(lines):
+                filtered_documents[index] = "\n".join(lines)
         # Return the filtered results instead of all results
         return {
             "query_embedding": query_embedding,
@@ -444,50 +489,52 @@ class ChromaDBFileWatcher(FileSystemEventHandler):
             "distances": filtered_distances,
             "metadatas": filtered_metadatas,
         }
-    
+
     def _get_file_hash(self, file_path):
         """Calculate MD5 hash of a file."""
         try:
-            with open(file_path, 'rb') as f:
+            with open(file_path, "rb") as f:
                 return hashlib.md5(f.read()).hexdigest()
         except Exception as e:
             logging.error(f"Error hashing file {file_path}: {e}")
             return None
-    
+
     def on_modified(self, event):
         """Handle file modification events."""
         if event.is_directory:
             return
-        
+
         file_path = event.src_path
         # Schedule the update using asyncio
         asyncio.run_coroutine_threadsafe(self.process_file_update(file_path), self.loop)
         logging.info(f"File modified: {file_path}")
-    
+
     def on_created(self, event):
         """Handle file creation events."""
         if event.is_directory:
             return
-        
+
         file_path = event.src_path
         # Schedule the update using asyncio
         asyncio.run_coroutine_threadsafe(self.process_file_update(file_path), self.loop)
         logging.info(f"File created: {file_path}")
-    
+
     def on_deleted(self, event):
         """Handle file deletion events."""
         if event.is_directory:
             return
-        
+
         file_path = event.src_path
-        asyncio.run_coroutine_threadsafe(self.remove_file_from_collection(file_path), self.loop)
+        asyncio.run_coroutine_threadsafe(
+            self.remove_file_from_collection(file_path), self.loop
+        )
         logging.info(f"File deleted: {file_path}")
 
     def on_moved(self, event):
         """Handle move deletion events."""
         if event.is_directory:
             return
-        
+
         file_path = event.src_path
         logging.info(f"TODO: on_moved: ${file_path}")
 
@@ -508,72 +555,88 @@ class ChromaDBFileWatcher(FileSystemEventHandler):
         try:
             # Remove existing entries for this file
             existing_results = self.collection.get(where={"path": file_path})
-            if existing_results and 'ids' in existing_results and existing_results['ids']:
-                self.collection.delete(ids=existing_results['ids'])
-            
+            if (
+                existing_results
+                and "ids" in existing_results
+                and existing_results["ids"]
+            ):
+                self.collection.delete(ids=existing_results["ids"])
+
             extensions = (".docx", ".xlsx", ".xls", ".pdf")
             if file_path.endswith(extensions):
                 p = Path(file_path)
                 p_as_md = p.with_suffix(".md")
                 if p_as_md.exists():
-                    logging.info(f"newer: {p.stat().st_mtime > p_as_md.stat().st_mtime}")
+                    logging.info(
+                        f"newer: {p.stat().st_mtime > p_as_md.stat().st_mtime}"
+                    )
 
                 # If file_path.md doesn't exist or file_path is newer than file_path.md,
                 # fire off markitdown
-                if (not p_as_md.exists()) or (p.stat().st_mtime > p_as_md.stat().st_mtime):
+                if (not p_as_md.exists()) or (
+                    p.stat().st_mtime > p_as_md.stat().st_mtime
+                ):
                     self._markitdown(file_path, p_as_md)
                 return
 
-            # Create document object in LangChain format
-            with open(file_path, "r", encoding="utf-8") as f:
-                content = f.read()
-            
+            chunks = self._markdown_chunker.process_file(file_path)
+            if not chunks:
+                return
+
             # Extract top-level directory
             rel_path = os.path.relpath(file_path, self.watch_directory)
-            top_level_dir = rel_path.split(os.sep)[0]
-            
-            document = Document(
-                page_content=content,
-                metadata={"doc_type": top_level_dir, "path": file_path}
-            )
-            
-            # Create chunks
-            chunks = self.text_splitter.split_documents([document])
-            
+            path_parts = rel_path.split(os.sep)
+            top_level_dir = path_parts[0]
+            # file_name = path_parts[-1]
+            for i, chunk in enumerate(chunks):
+                chunk["metadata"]["doc_type"] = top_level_dir
+                # with open(f"src/tmp/{file_name}.{i}", "w") as f:
+                #     f.write(json.dumps(chunk, indent=2))
+
             # Add chunks to collection
             self.add_embeddings_to_collection(chunks)
-            
+
             logging.info(f"Updated {len(chunks)} chunks for file: {file_path}")
-            
+
         except Exception as e:
             logging.error(f"Error updating document in collection: {e}")
-    
+            logging.error(traceback.format_exc())
+
     async def initialize_collection(self):
         """Initialize the collection with all documents from the watch directory."""
         # Process all files regardless of hash state
         num_processed = await self.scan_directory(process_all=True)
-        
-        logging.info(f"Vectorstore initialized with {self.collection.count()} documents")
-        
+
+        logging.info(
+            f"Vectorstore initialized with {self.collection.count()} documents"
+        )
+
         self._update_umaps()
 
         # Show stats
         try:
-            all_metadata = self.collection.get()['metadatas']
+            all_metadata = self.collection.get()["metadatas"]
             if all_metadata:
-                doc_types = set(m.get('doc_type', 'unknown') for m in all_metadata)
+                doc_types = set(m.get("doc_type", "unknown") for m in all_metadata)
                 logging.info(f"Document types: {doc_types}")
         except Exception as e:
             logging.error(f"Error getting document types: {e}")
-            
+
         return num_processed
 
+
 # Function to start the file watcher
-def start_file_watcher(llm, watch_directory, persist_directory=None, 
-                      collection_name="documents", initialize=False, recreate=False):
+def start_file_watcher(
+    llm,
+    watch_directory,
+    persist_directory=None,
+    collection_name="documents",
+    initialize=False,
+    recreate=False,
+):
     """
     Start watching a directory for file changes.
-    
+
     Args:
         llm: The language model client
         watch_directory: Directory to watch for changes
@@ -585,14 +648,14 @@ def start_file_watcher(llm, watch_directory, persist_directory=None,
     loop = asyncio.get_event_loop()
 
     file_watcher = ChromaDBFileWatcher(
-        llm, 
-        watch_directory, 
+        llm,
+        watch_directory,
         loop=loop,
         persist_directory=persist_directory,
         collection_name=collection_name,
-        recreate=recreate
+        recreate=recreate,
     )
-    
+
     # Process all files if:
     # 1. initialize=True was passed (explicit request to initialize)
     # 2. This is a new collection (doesn't exist yet)
@@ -604,38 +667,39 @@ def start_file_watcher(llm, watch_directory, persist_directory=None,
         # Only process new/changed files
         logging.info("Scanning for new/changed documents")
         asyncio.run_coroutine_threadsafe(file_watcher.scan_directory(), loop)
-    
+
     # Start observer
     observer = Observer()
     observer.schedule(file_watcher, watch_directory, recursive=True)
     observer.start()
-    
+
     logging.info(f"Started watching directory: {watch_directory}")
     return observer, file_watcher
 
+
 if __name__ == "__main__":
     # When running directly, use absolute imports
     import defines
-    
+
     # Initialize Ollama client
-    llm = ollama.Client(host=defines.ollama_api_url) # type: ignore
-    
+    llm = ollama.Client(host=defines.ollama_api_url)  # type: ignore
+
     # Start the file watcher (with initialization)
     observer, file_watcher = start_file_watcher(
-        llm, 
-        defines.doc_dir, 
+        llm,
+        defines.doc_dir,
         recreate=True,  # Start fresh
     )
-    
+
     # Example query
     query = "Can you describe James Ketrenos' work history?"
     top_docs = file_watcher.find_similar(query, top_k=3)
     logging.info(top_docs)
-    
+
     try:
         # Keep the main thread running
         while True:
             time.sleep(1)
     except KeyboardInterrupt:
         observer.stop()
-    observer.join()
\ No newline at end of file
+    observer.join()
diff --git a/src/utils/setup_logging.py b/src/utils/setup_logging.py
index d2c6655..c1c85b4 100644
--- a/src/utils/setup_logging.py
+++ b/src/utils/setup_logging.py
@@ -4,9 +4,12 @@ import logging
 
 from . import defines
 
+
 def setup_logging(level=defines.logging_level) -> logging.Logger:
     os.environ["TORCH_CPP_LOG_LEVEL"] = "ERROR"
-    warnings.filterwarnings("ignore", message="Overriding a previously registered kernel")
+    warnings.filterwarnings(
+        "ignore", message="Overriding a previously registered kernel"
+    )
     warnings.filterwarnings("ignore", message="Warning only once for all operators")
     warnings.filterwarnings("ignore", message=".*Couldn't find ffmpeg or avconv.*")
     warnings.filterwarnings("ignore", message="'force_all_finite' was renamed to")
@@ -16,19 +19,25 @@ def setup_logging(level=defines.logging_level) -> logging.Logger:
     numeric_level = getattr(logging, level.upper(), None)
     if not isinstance(numeric_level, int):
         raise ValueError(f"Invalid log level: {level}")
-    
+
     logging.basicConfig(
-        level=numeric_level, 
+        level=numeric_level,
         format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s",
         datefmt="%Y-%m-%d %H:%M:%S",
-        force=True
+        force=True,
     )
 
     # Now reduce verbosity for FastAPI, Uvicorn, Starlette
-    for noisy_logger in ("uvicorn", "uvicorn.error", "uvicorn.access", "fastapi", "starlette"):
-    #for noisy_logger in ("starlette"):
+    for noisy_logger in (
+        "uvicorn",
+        "uvicorn.error",
+        "uvicorn.access",
+        "fastapi",
+        "starlette",
+    ):
+        # for noisy_logger in ("starlette"):
         logging.getLogger(noisy_logger).setLevel(logging.WARNING)
-    
+
     logger = logging.getLogger(__name__)
 
     return logger
diff --git a/src/utils/tools/__init__.py b/src/utils/tools/__init__.py
index ba2a83f..c03439c 100644
--- a/src/utils/tools/__init__.py
+++ b/src/utils/tools/__init__.py
@@ -1,14 +1,14 @@
 import importlib
 
-from . basetools import tools, llm_tools, enabled_tools, tool_functions
-from .. setup_logging import setup_logging
+from .basetools import tools, llm_tools, enabled_tools, tool_functions
+from ..setup_logging import setup_logging
 from .. import defines
 
 logger = setup_logging(level=defines.logging_level)
 
 # Dynamically import all names from basetools listed in tools_all
-module = importlib.import_module('.basetools', package=__package__)
+module = importlib.import_module(".basetools", package=__package__)
 for name in tool_functions:
     globals()[name] = getattr(module, name)
 
-__all__ = [ 'tools', 'llm_tools', 'enabled_tools', 'tool_functions' ]
+__all__ = ["tools", "llm_tools", "enabled_tools", "tool_functions"]
diff --git a/src/utils/tools/basetools.py b/src/utils/tools/basetools.py
index 30d7e56..7044a75 100644
--- a/src/utils/tools/basetools.py
+++ b/src/utils/tools/basetools.py
@@ -1,18 +1,19 @@
 import os
 from datetime import datetime
 from typing import (
-  Any,
+    Any,
 )
 from typing_extensions import Annotated
 
-from bs4 import BeautifulSoup # type: ignore
+from bs4 import BeautifulSoup  # type: ignore
 
-from geopy.geocoders import Nominatim # type: ignore
-import pytz # type: ignore
+from geopy.geocoders import Nominatim  # type: ignore
+import pytz  # type: ignore
 import requests
-import yfinance as yf # type: ignore
+import yfinance as yf  # type: ignore
 import logging
 
+
 # %%
 def WeatherForecast(city, state, country="USA"):
     """
@@ -42,24 +43,25 @@ def WeatherForecast(city, state, country="USA"):
     # Step 3: Get the forecast data from the grid endpoint
     forecast = get_forecast(grid_endpoint)
 
-    if not forecast['location']:
-        forecast['location'] = location
+    if not forecast["location"]:
+        forecast["location"] = location
 
     return forecast
 
+
 def get_coordinates(location):
     """Convert a location string to latitude and longitude using Nominatim geocoder."""
     try:
         # Create a geocoder with a meaningful user agent
         geolocator = Nominatim(user_agent="weather_app_example")
-        
+
         # Get the location
         location_data = geolocator.geocode(location)
-        
+
         if location_data:
             return {
                 "latitude": location_data.latitude,
-                "longitude": location_data.longitude
+                "longitude": location_data.longitude,
             }
         else:
             print(f"Location not found: {location}")
@@ -68,22 +70,23 @@ def get_coordinates(location):
         print(f"Error getting coordinates: {e}")
         return None
 
+
 def get_grid_endpoint(coordinates):
     """Get the grid endpoint from weather.gov based on coordinates."""
     try:
         lat = coordinates["latitude"]
         lon = coordinates["longitude"]
-        
+
         # Define headers for the API request
         headers = {
             "User-Agent": "WeatherAppExample/1.0 (your_email@example.com)",
-            "Accept": "application/geo+json"
+            "Accept": "application/geo+json",
         }
-        
+
         # Make the request to get the grid endpoint
         url = f"https://api.weather.gov/points/{lat},{lon}"
         response = requests.get(url, headers=headers)
-        
+
         if response.status_code == 200:
             data = response.json()
             return data["properties"]["forecast"]
@@ -94,44 +97,50 @@ def get_grid_endpoint(coordinates):
         print(f"Error in get_grid_endpoint: {e}")
         return None
 
+
 # Weather related function
 
+
 def get_forecast(grid_endpoint):
     """Get the forecast data from the grid endpoint."""
     try:
         # Define headers for the API request
         headers = {
             "User-Agent": "WeatherAppExample/1.0 (your_email@example.com)",
-            "Accept": "application/geo+json"
+            "Accept": "application/geo+json",
         }
-        
+
         # Make the request to get the forecast
         response = requests.get(grid_endpoint, headers=headers)
-        
+
         if response.status_code == 200:
             data = response.json()
-            
+
             # Extract the relevant forecast information
             periods = data["properties"]["periods"]
-            
+
             # Process the forecast data into a simpler format
             forecast = {
-                "location": data["properties"].get("relativeLocation", {}).get("properties", {}),
+                "location": data["properties"]
+                .get("relativeLocation", {})
+                .get("properties", {}),
                 "updated": data["properties"].get("updated", ""),
-                "periods": []
+                "periods": [],
             }
-            
+
             for period in periods:
-                forecast["periods"].append({
-                    "name": period.get("name", ""),
-                    "temperature": period.get("temperature", ""),
-                    "temperatureUnit": period.get("temperatureUnit", ""),
-                    "windSpeed": period.get("windSpeed", ""),
-                    "windDirection": period.get("windDirection", ""),
-                    "shortForecast": period.get("shortForecast", ""),
-                    "detailedForecast": period.get("detailedForecast", "")
-                })
-            
+                forecast["periods"].append(
+                    {
+                        "name": period.get("name", ""),
+                        "temperature": period.get("temperature", ""),
+                        "temperatureUnit": period.get("temperatureUnit", ""),
+                        "windSpeed": period.get("windSpeed", ""),
+                        "windDirection": period.get("windDirection", ""),
+                        "shortForecast": period.get("shortForecast", ""),
+                        "detailedForecast": period.get("detailedForecast", ""),
+                    }
+                )
+
             return forecast
         else:
             print(f"Error getting forecast: {response.status_code} - {response.text}")
@@ -140,15 +149,16 @@ def get_forecast(grid_endpoint):
         print(f"Error in get_forecast: {e}")
         return {"error": f"Exception: {str(e)}"}
 
+
 # Example usage
 # def do_weather():
 #     city = input("Enter city: ")
 #     state = input("Enter state: ")
 #     country = input("Enter country (default USA): ") or "USA"
-    
+
 #     print(f"Getting weather for {city}, {state}, {country}...")
 #     weather_data = WeatherForecast(city, state, country)
-    
+
 #     if "error" in weather_data:
 #         print(f"Error: {weather_data['error']}")
 #     else:
@@ -156,7 +166,7 @@ def get_forecast(grid_endpoint):
 #         print(f"Location: {weather_data.get('location', {}).get('city', city)}, {weather_data.get('location', {}).get('state', state)}")
 #         print(f"Last Updated: {weather_data.get('updated', 'N/A')}")
 #         print("\nForecast Periods:")
-        
+
 #         for period in weather_data.get("periods", []):
 #             print(f"\n{period['name']}:")
 #             print(f"  Temperature: {period['temperature']}{period['temperatureUnit']}")
@@ -166,35 +176,32 @@ def get_forecast(grid_endpoint):
 
 # %%
 
+
 def TickerValue(ticker_symbols):
     api_key = os.getenv("TWELVEDATA_API_KEY", "")
     if not api_key:
         return {"error": f"Error fetching data: No API key for TwelveData"}
-    
+
     results = []
-    for ticker_symbol in ticker_symbols.split(','):
+    for ticker_symbol in ticker_symbols.split(","):
         ticker_symbol = ticker_symbol.strip()
         if ticker_symbol == "":
             continue
 
-        url = f"https://api.twelvedata.com/price?symbol={ticker_symbol}&apikey={api_key}"
-    
+        url = (
+            f"https://api.twelvedata.com/price?symbol={ticker_symbol}&apikey={api_key}"
+        )
+
         response = requests.get(url)
         data = response.json()
 
         if "price" in data:
             logging.info(f"TwelveData: {ticker_symbol} {data}")
-            results.append({
-                "symbol": ticker_symbol,
-                "price": float(data["price"])
-            })
+            results.append({"symbol": ticker_symbol, "price": float(data["price"])})
         else:
             logging.error(f"TwelveData: {data}")
-            results.append({
-                "symbol": ticker_symbol,
-                "price": "Unavailable"
-            })
-        
+            results.append({"symbol": ticker_symbol, "price": "Unavailable"})
+
     return results[0] if len(results) == 1 else results
 
 
@@ -202,15 +209,15 @@ def TickerValue(ticker_symbols):
 def yfTickerValue(ticker_symbols):
     """
     Look up the current price of a stock using its ticker symbol.
-    
+
     Args:
         ticker_symbol (str): The stock ticker symbol (e.g., 'AAPL' for Apple)
-        
+
     Returns:
         dict: Current stock information including price
     """
     results = []
-    for ticker_symbol in ticker_symbols.split(','):
+    for ticker_symbol in ticker_symbols.split(","):
         ticker_symbol = ticker_symbol.strip()
         if ticker_symbol == "":
             continue
@@ -220,59 +227,64 @@ def yfTickerValue(ticker_symbols):
             ticker = yf.Ticker(ticker_symbol)
             # Get the latest market data
             ticker_data = ticker.history(period="1d")
-            
+
             if ticker_data.empty:
                 results.append({"error": f"No data found for ticker {ticker_symbol}"})
                 continue
-            
+
             # Get the latest closing price
-            latest_price = ticker_data['Close'].iloc[-1]
-            
+            latest_price = ticker_data["Close"].iloc[-1]
+
             # Get some additional info
-            results.append({ 'symbol': ticker_symbol, 'price': latest_price })
-    
+            results.append({"symbol": ticker_symbol, "price": latest_price})
+
         except Exception as e:
             import traceback
+
             logging.error(f"Error fetching data for {ticker_symbol}: {e}")
             logging.error(traceback.format_exc())
-            results.append({"error": f"Error fetching data for {ticker_symbol}: {str(e)}"})
-            
+            results.append(
+                {"error": f"Error fetching data for {ticker_symbol}: {str(e)}"}
+            )
+
     return results[0] if len(results) == 1 else results
 
+
 # %%
 def DateTime(timezone="America/Los_Angeles"):
     """
     Returns the current date and time in the specified timezone in ISO 8601 format.
-    
+
     Args:
         timezone (str): Timezone name (e.g., "UTC", "America/New_York", "Europe/London")
                         Default is "America/Los_Angeles"
-    
+
     Returns:
         str: Current date and time with timezone in the format YYYY-MM-DDTHH:MM:SS+HH:MM
     """
     try:
-        if timezone == 'system' or timezone == '' or not timezone:
-            timezone = 'America/Los_Angeles'
+        if timezone == "system" or timezone == "" or not timezone:
+            timezone = "America/Los_Angeles"
         # Get current UTC time (timezone-aware)
         local_tz = pytz.timezone("America/Los_Angeles")
         local_now = datetime.now(tz=local_tz)
-        
+
         # Convert to target timezone
         target_tz = pytz.timezone(timezone)
         target_time = local_now.astimezone(target_tz)
-        
+
         return target_time.isoformat()
     except Exception as e:
-        return {'error': f"Invalid timezone {timezone}: {str(e)}"}
-        
-async def AnalyzeSite(llm, model: str, url : str, question : str):
+        return {"error": f"Invalid timezone {timezone}: {str(e)}"}
+
+
+async def AnalyzeSite(llm, model: str, url: str, question: str):
     """
     Fetches content from a URL, extracts the text, and uses Ollama to summarize it.
-    
+
     Args:
         url (str): The URL of the website to summarize
-        
+
     Returns:
         str: A summary of the website content
     """
@@ -287,41 +299,43 @@ async def AnalyzeSite(llm, model: str, url : str, question : str):
         logging.info(f"{url} returned. Processing...")
         # Parse the HTML
         soup = BeautifulSoup(response.text, "html.parser")
-        
+
         # Remove script and style elements
         for script in soup(["script", "style"]):
             script.extract()
-            
+
         # Get text content
         text = soup.get_text(separator=" ", strip=True)
-        
+
         # Clean up text (remove extra whitespace)
         lines = (line.strip() for line in text.splitlines())
         chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
         text = " ".join(chunk for chunk in chunks if chunk)
-        
+
         # Limit text length if needed (Ollama may have token limits)
         max_chars = 100000
         if len(text) > max_chars:
             text = text[:max_chars] + "..."
-            
+
         # Create Ollama client
         # logging.info(f"Requesting summary of: {text}")
-        
+
         # Generate summary using Ollama
         prompt = f"CONTENTS:\n\n{text}\n\n{question}"
-        response = llm.generate(model=model, 
-                                   system="You are given the contents of {url}. Answer the question about the contents",
-                                   prompt=prompt)
-        
-        #logging.info(response["response"])
-        
+        response = llm.generate(
+            model=model,
+            system="You are given the contents of {url}. Answer the question about the contents",
+            prompt=prompt,
+        )
+
+        # logging.info(response["response"])
+
         return {
             "source": "summarizer-llm",
             "content": response["response"],
-            "metadata": DateTime()
+            "metadata": DateTime(),
         }
-        
+
     except requests.exceptions.RequestException as e:
         logging.error(f"Error fetching the URL: {e}")
         return f"Error fetching the URL: {str(e)}"
@@ -331,109 +345,116 @@ async def AnalyzeSite(llm, model: str, url : str, question : str):
 
 
 # %%
-tools = [ {
-    "type": "function", 
-    "function": {
-        "name": "TickerValue",
-        "description": "Get the current stock price of one or more ticker symbols. Returns an array of objects with 'symbol' and 'price' fields. Call this whenever you need to know the latest value of stock ticker symbols, for example when a user asks 'How much is Intel trading at?' or 'What are the prices of AAPL and MSFT?'",
-        "parameters": {
-            "type": "object",
-            "properties": {
-                "ticker": {
-                    "type": "string",
-                    "description": "The company stock ticker symbol. For multiple tickers, provide a comma-separated list (e.g., 'AAPL,MSFT,GOOGL').",
+tools = [
+    {
+        "type": "function",
+        "function": {
+            "name": "TickerValue",
+            "description": "Get the current stock price of one or more ticker symbols. Returns an array of objects with 'symbol' and 'price' fields. Call this whenever you need to know the latest value of stock ticker symbols, for example when a user asks 'How much is Intel trading at?' or 'What are the prices of AAPL and MSFT?'",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "ticker": {
+                        "type": "string",
+                        "description": "The company stock ticker symbol. For multiple tickers, provide a comma-separated list (e.g., 'AAPL,MSFT,GOOGL').",
+                    },
                 },
+                "required": ["ticker"],
+                "additionalProperties": False,
             },
-            "required": ["ticker"],
-            "additionalProperties": False
-        }
-    }
-}, {
-    "type": "function", 
-    "function": {
-        "name": "AnalyzeSite",
-        "description": "Downloads the requested site and asks a second LLM agent to answer the question based on the site content. For example if the user says 'What are the top headlines on cnn.com?' you would use AnalyzeSite to get the answer. Only use this if the user asks about a specific site or company.",
-        "parameters": {
-            "type": "object",
-            "properties": {
-                "url": {
-                    "type": "string",
-                    "description": "The website URL to download and process",
-                },
-                "question": {
-                    "type": "string",
-                    "description": "The question to ask the second LLM about the content",
-                },
-            },
-            "required": ["url", "question"],
-            "additionalProperties": False
         },
-        "returns": {
-            "type": "object",
-            "properties": {
-                "source": {
-                    "type": "string",
-                    "description": "Identifier for the source LLM"
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "AnalyzeSite",
+            "description": "Downloads the requested site and asks a second LLM agent to answer the question based on the site content. For example if the user says 'What are the top headlines on cnn.com?' you would use AnalyzeSite to get the answer. Only use this if the user asks about a specific site or company.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "url": {
+                        "type": "string",
+                        "description": "The website URL to download and process",
+                    },
+                    "question": {
+                        "type": "string",
+                        "description": "The question to ask the second LLM about the content",
+                    },
                 },
-                "content": {
-                    "type": "string",
-                    "description": "The complete response from the second LLM"
-                },
-                "metadata": {
-                    "type": "object",
-                    "description": "Additional information about the response"
-                }
-            }
-        }
-    }
-}, {
-    "type": "function", 
-    "function": {
-        "name": "DateTime",
-        "description": "Get the current date and time in a specified timezone. For example if a user asks 'What time is it in Poland?' you would pass the Warsaw timezone to DateTime.",
-        "parameters": {
-            "type": "object",
-            "properties": {
-                "timezone": {
-                    "type": "string",
-                    "description": "Timezone name (e.g., 'UTC', 'America/New_York', 'Europe/London', 'America/Los_Angeles'). Default is 'America/Los_Angeles'."
-                }
+                "required": ["url", "question"],
+                "additionalProperties": False,
             },
-            "required": []
-        }
-    }
-}, {
-    "type": "function", 
-    "function": {
-        "name": "WeatherForecast",
-        "description": "Get the full weather forecast as structured data for a given CITY and STATE location in the United States. For example, if the user asks 'What is the weather in Portland?' or 'What is the forecast for tomorrow?' use the provided data to answer the question.",
-        "parameters": {
-            "type": "object",
-            "properties": {
-                "city": {
-                    "type": "string",
-                    "description": "City to find the weather forecast (e.g., 'Portland', 'Seattle').",
-                    "minLength": 2
+            "returns": {
+                "type": "object",
+                "properties": {
+                    "source": {
+                        "type": "string",
+                        "description": "Identifier for the source LLM",
+                    },
+                    "content": {
+                        "type": "string",
+                        "description": "The complete response from the second LLM",
+                    },
+                    "metadata": {
+                        "type": "object",
+                        "description": "Additional information about the response",
+                    },
                 },
-                "state": {
-                    "type": "string",
-                    "description": "State to find the weather forecast (e.g., 'OR', 'WA').",
-                    "minLength": 2
-                }
             },
-            "required": [ "city", "state" ],
-            "additionalProperties": False
-        }
-    }
-}]
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "DateTime",
+            "description": "Get the current date and time in a specified timezone. For example if a user asks 'What time is it in Poland?' you would pass the Warsaw timezone to DateTime.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "timezone": {
+                        "type": "string",
+                        "description": "Timezone name (e.g., 'UTC', 'America/New_York', 'Europe/London', 'America/Los_Angeles'). Default is 'America/Los_Angeles'.",
+                    }
+                },
+                "required": [],
+            },
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "WeatherForecast",
+            "description": "Get the full weather forecast as structured data for a given CITY and STATE location in the United States. For example, if the user asks 'What is the weather in Portland?' or 'What is the forecast for tomorrow?' use the provided data to answer the question.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "city": {
+                        "type": "string",
+                        "description": "City to find the weather forecast (e.g., 'Portland', 'Seattle').",
+                        "minLength": 2,
+                    },
+                    "state": {
+                        "type": "string",
+                        "description": "State to find the weather forecast (e.g., 'OR', 'WA').",
+                        "minLength": 2,
+                    },
+                },
+                "required": ["city", "state"],
+                "additionalProperties": False,
+            },
+        },
+    },
+]
+
 
 def llm_tools(tools):
     return [tool for tool in tools if tool.get("enabled", False) == True]
 
+
 def enabled_tools(tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
     return [{**tool, "enabled": True} for tool in tools]
 
-tool_functions = [ 'DateTime', 'WeatherForecast', 'TickerValue', 'AnalyzeSite' ]
-__all__ = [ 'tools', 'llm_tools', 'enabled_tools', 'tool_functions' ]
-#__all__.extend(__tool_functions__) # type: ignore
 
+tool_functions = ["DateTime", "WeatherForecast", "TickerValue", "AnalyzeSite"]
+__all__ = ["tools", "llm_tools", "enabled_tools", "tool_functions"]
+# __all__.extend(__tool_functions__) # type: ignore