Update tool_handler.py

tool_handler.py

@@ -51,18 +51,28 @@ def chat_memory_tool(memory: str, model="llama-3.3-70b-versatile", use_rag=True)
 def treatment_tool(query: str, model="llama-3.3-70b-versatile", use_rag=True) -> str:
     try:
         query_embedding = get_embedding(query)
-        results = discharge_collection.query(
-            query_embeddings=[query_embedding],
-            n_results=5,
-            include=["documents"]
-        )
-        top_docs = results['documents'][0] if use_rag else []
-        combined_context = "\n\n".join(top_docs) if use_rag else ""
-        prompt = (
-            "You are a helpful medical assistant. Based on the following discharge notes, recommend essential treatment.\n\n"
-            f"### Notes:\n{combined_context}\n\n### Condition:\n{query}"
-            if use_rag else f"Patient condition: {query}. What treatment is recommended?"
-        )
+
+        # If RAG is on, query vector DB and prepare context
+        if use_rag:
+            results = discharge_collection.query(
+                query_embeddings=[query_embedding],
+                n_results=5,
+                include=["documents"]
+            )
+
+            top_docs = results['documents'][0] if results and results['documents'] else []
+            
+            # Truncate each doc to avoid token overflow
+            top_docs = [doc[:1500] for doc in top_docs]
+            combined_context = "\n\n".join(top_docs)
+
+            prompt = (
+                "You are a helpful medical assistant. Based on the following discharge notes, recommend essential treatment.\n\n"
+                f"### Notes:\n{combined_context}\n\n### Condition:\n{query}"
+            )
+        else:
+            prompt = f"Patient condition: {query}. What treatment is recommended?"
+
         response = client.chat.completions.create(
             model=model,
             messages=[
@@ -71,6 +81,7 @@ def treatment_tool(query: str, model="llama-3.3-70b-versatile", use_rag=True) ->
             ]
         )
         return response.choices[0].message.content
+
     except Exception as e:
         return f"Error: {str(e)}"