Upload 26 files

healthagent/agents/analyzer.py

@@ -0,0 +1,41 @@
+from typing import Dict, Any
+from pydantic import BaseModel, Field
+from healthagent.predictive.custom_model import ModelAPI
+from healthagent.tools.explainer import XiaExplainer
+
+modelAPI = ModelAPI()
+xiaAPI = XiaExplainer()
+
+class PredictionInput(BaseModel):
+    """Schema for prediction input parameters."""
+    driver_temperature: float = Field(..., description="Driver temperature in Celsius")
+    temperature: float = Field(..., description="Temperature in Celsius")
+    speed_rpm: float = Field(..., description="Speed in RPM")
+    current: float = Field(..., description="Current in Amperes")
+    voltage: float = Field(..., description="Voltage in Volts")
+    power: float = Field(..., description="Power in Watts")
+    thruster_id_encoded: float = Field(default=1.0, description="Encoded thruster ID (default=1.0)")
+
+
+async def prediction_causes_tool(input_data: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Useful for identifying the causes of the model predictions.
+
+    Args:
+        input_data  (Dict[str, Any]): Raw input values for prediction.
+
+    Returns:
+        Dict[str, Any]: Mapping of feature -> value.
+    """
+    
+    input_obj = PredictionInput(**input_data)
+    input_dict = input_obj.model_dump()
+    input_dict["thruster_id_encoded"] = "1.0" 
+     
+    result: List[Dict[str, str]]  = xiaAPI.explain_prediction(modelAPI.model, input_dict, True)
+
+    formatted_result: Dict[str, Any] = dict(
+        map(lambda item: (item["feature"], str(item["value"])), result)
+    )
+     
+    return  formatted_result

healthagent/agents/summarization.py

@@ -0,0 +1,24 @@
+from llama_index.core import Settings
+from llama_index.core import PromptTemplate
+from typing import Dict, Any
+
+async def writing_tool(input_data: Dict[str, Any]) -> str:
+    """
+    Useful for writing a summary of the given context.
+
+    """
+    PROMPT = PromptTemplate(
+    template="""
+        Write a detailed information of the given context.
+        The context contains information about causes of prediction, the state of the threshold condtions and how to troubleshoot.
+        
+
+        Here are the context:
+        {component}
+        """
+    )
+    query = PROMPT.format(component=input_data)
+     
+    #TODO: Verify Content Loop 
+    response =   Settings.llm.complete(query) 
+    return  response

healthagent/agents/threshold.py

@@ -0,0 +1,19 @@
+async def  threshold_tool(causes: str) -> str:
+  """Useful for writing the threshold information of the causes.""" 
+  """For each cause state whether the value of the causes are greater than the high threshold values or lesser than the threshold values only."""
+  """The threshold values given below. Return only the threshold values whose causes were identified """
+
+  threshold = {
+        "voltage": {"low": 20, "high": 28},
+        "current": {"low": 0, "high": 15},
+        "temperature": {"low": 0, "high": 75},
+        "driver_temperature": {"low": 0, "high": 80},
+        "speed": {"low": 1000, "high": 5000},
+        "power": {"low": 200, "high": 1000},
+    }
+
+   
+  #TODO: Select only the influencing threshold values
+  
+  return str(threshold)
+ 

healthagent/config/model_config.yaml

@@ -0,0 +1,17 @@
+model:
+  language_name: qwen2.5:7b
+  embedding: bge-m3
+  predictive_model: healthagent/model/predictive_model_RandomForest.pkl
+  xia_model: healthagent/model/xia_config.pkl
+
+chunk:
+  size: 512
+  overlap: 64
+
+url:
+  server: http://api:8080/
+
+documents:
+  breakpoints: ./healthagent/data/documents/
+
+storage: ./healthagent/storage/

healthagent/predictive/analysis.py

@@ -0,0 +1,8 @@
+import random
+## This function is Deprecated and will be removed in future version
+def get_status(data=None, model=None) -> str:
+
+    STATUS_LABELS = ['FAULTY', 'NORMAL', 'WARNING']
+    prediction = model.predict(data)[0]
+
+    return STATUS_LABELS[int(prediction.item())]

healthagent/predictive/custom_model.py

@@ -0,0 +1,20 @@
+import joblib
+from healthagent.utils.setup import PREDICTIVE_MODEL
+from healthagent.utils.data_lib import data_seralize
+from healthagent.utils.logger import setup_logger
+
+logger = setup_logger()
+STATUS_LABEL = ['FAULTY', 'NORMAL', 'WARNING'] #TODO: joblib for the label
+
+class ModelAPI:
+  def __init__(self):
+    logger.info("Configuring predictive model.")
+    self.model = joblib.load(PREDICTIVE_MODEL)
+    self.STATUS_LABEL = STATUS_LABEL
+    logger.info("Predictive model configured successfully.")
+
+  def predict(self, x):
+    sample_input= data_seralize(x)
+    prediction = self.model.predict(sample_input)[0]
+    value = int(prediction.item())
+    return self.STATUS_LABEL[value]

healthagent/slm/context_state.py

@@ -0,0 +1,12 @@
+from llama_index.core.workflow import Context
+
+async def record_notes(ctx: Context, title:str, content: str) -> str:
+    """Useful for recording notes on a given results. Your content should be the causes of predictions."""
+
+    #TODO: Perform code struct
+    async with ctx.store.edit_state() as ctx_state:
+        if "_notes" not in ctx_state["state"]:
+            ctx_state["state"]["_notes"] = {}
+        ctx_state["state"]["_notes"][title] = content
+    
+    return "Notes recorded."

healthagent/slm/core.py

@@ -0,0 +1,24 @@
+from llama_index.core import Settings
+from llama_index.llms.ollama import Ollama
+from llama_index.embeddings.ollama import OllamaEmbedding
+
+from healthagent.utils.setup import MODEL_NAME, EMBED_MODEL, CHUNK_SIZE, CHUNK_OVERLAP
+from healthagent.utils.logger import setup_logger
+
+logger = setup_logger()
+
+def configure_llm():
+    logger.info("Setting up language model.")
+    Settings.llm = language_model()
+    Settings.embed_model = OllamaEmbedding(model_name=EMBED_MODEL,ollama_additional_kwargs={"mirostat": 0})
+    Settings.chunk_size = CHUNK_SIZE
+    Settings.chunk_overlap = CHUNK_OVERLAP
+    logger.info("Language model configured sucessfully.")
+
+
+def language_model():
+  return Ollama(model=MODEL_NAME,
+                request_timeout=360,
+                temperature=1.0,
+                context_window=8000
+                )

healthagent/slm/operation.py

@@ -0,0 +1,73 @@
+from llama_index.core.agent.workflow import FunctionAgent, AgentWorkflow
+from llama_index.core.tools import QueryEngineTool
+from healthagent.agents.analyzer import prediction_causes_tool
+from healthagent.agents.threshold import threshold_tool
+from llama_index.core.workflow import Context
+from healthagent.slm.context_state import record_notes, write_report
+from healthagent.agents.summarization import writing_tool
+from healthagent.slm.retrieval import Retrieval
+
+
+def causes_agents():
+    return FunctionAgent (
+        name="PredictionCausesAgent",
+        system_prompt=(
+            "You are the PredictionCausesAgent that provides the causes of the model prediction."
+            "Record note of the results with the title Causes of model predictions."
+        ),
+        tools=[prediction_causes_tool, record_notes],
+        description="Useful for identifying the causes of the model predictions and record notes on the results.",
+        can_handoff_to=["ThresholdAgent"],
+    )
+
+
+def threshold_agents():
+    return FunctionAgent(
+        name="ThresholdAgent",
+        system_prompt=(
+           "You are the ThresholdAgent that deterimine if the causes values are low or high "
+            "Record note of the results with the title troubleshoot information."
+            
+        ),
+        tools=[threshold_tool, record_notes],
+        description="Useful for deterimine if the causes values are low or high and record notes on the results with the title threshold information.",
+
+        can_handoff_to=["TroubleshootDocumentAgent"],
+    )
+
+
+
+
+def write_agent():
+    return  FunctionAgent(
+    name="ReportingAgent",
+    description="Useful for summaring all the above information as a a report.",
+    system_prompt=(
+        "You are the WriteAgent that summaries all the above information as a a report "
+        "Write your responses in the write note tool"
+        "Your report should be in a markdown format and use the write_report tool to compile."
+       
+    ),
+    tools=[writing_tool, write_report],
+    can_handoff_to=[]
+
+)
+
+
+def multi_workflow():
+    causesAgent = causes_agents()
+    thresholdAgent = threshold_agents()
+    
+    writeAgent = write_agent()
+
+
+    workflow = AgentWorkflow(
+        agents=[causesAgent, thresholdAgent,   writeAgent],
+        root_agent=causesAgent.name,
+        initial_state={
+            "_notes": {},
+            "_content": "Not written yet.",
+        },
+
+    )
+    return workflow

healthagent/slm/prompts.py

@@ -0,0 +1,41 @@
+from llama_index.core import PromptTemplate
+from typing import List, Dict
+
+#Deprecated: This function will be removed in future version.
+def analysis_prompt(data: Dict[str, str], context_info) -> str:
+
+    PROMPT = PromptTemplate(
+    template="""
+        You are a Remote Operating Vehicle (ROV) assistant. Based on the sensor data provided, write a brief summary sentence of the component's health condition based on the status and the causes provided.
+
+        Your response should:
+            Begin the sentence with the component name.
+            Be brief and simple.
+            Indicate the status.
+            Mention any potential problem if sensor readings indicate a possible issue based on the causes.
+            Do not prefix the result with any labels like "Summary:".
+
+        Here are the information:
+        Component Name: {component}
+        Driver Temperature: {driver_temperature} in degrees Celsius
+        Temperature: {temperature} in degrees Celsius
+        Speed: {speed} in RPM
+        Current: {current} in Amperes
+        Power: {power}
+        Voltage: {voltage}
+        Status: {status}
+        Causes: {causes}
+        """)
+
+    return PROMPT.format(
+        component=data["component"],
+        driver_temperature=data["driver_temperature"],
+        temperature=data["temperature"],
+        speed=data["speed_rpm"],
+        current=data["current"],
+        power=data["power"],
+        voltage=data["voltage"],
+        status=data["status"],
+        causes=data["causes"],
+        context=context_info
+    )

healthagent/slm/query.py

@@ -0,0 +1,35 @@
+from healthagent.types import ThrusterData
+from llama_index.core.prompts import RichPromptTemplate
+
+
+def get_cause_query(thruster_data):
+  template_str = "Identify the cause of the model prediction with the following information: {{ thrusterData }}\n"
+  prompt = RichPromptTemplate(template_str)
+  data = ThrusterData(**thruster_data)
+  
+  return prompt.format(thrusterData=data)
+
+#TODO: Deprecated function. Remove this in next version
+def xia_query(data):
+  voltage = data["voltage"]
+  current = data["current"]
+  power = data["power"]
+  temperature = data["temperature"]
+  driver_temperature = data["driver_temperature"]
+  speed_rpm = data["speed_rpm"]
+  thruster_id = data["thruster_id_encoded"]
+
+  query = f"Identify and record the causes of the model predictions with driver_temperature of {driver_temperature}, current of {current}, speed_rpm of {speed_rpm}, voltage of {voltage}, temperature of {temperature}, power of {power}."
+  return query
+ 
+
+def support_query():
+  threshold = "Identify and record if the model predictions causes values are low or high."    
+  troubleshoot = "Write and record the troubleshooting steps for the threshold results."
+  report = "Summaries all the above information as a a report"
+    
+  return  f"""
+  {threshold}
+  {troubleshoot}
+  {report}
+  """

healthagent/slm/retrieval.py

@@ -0,0 +1,114 @@
+import os
+from llama_index.core.node_parser import SemanticSplitterNodeParser, SentenceSplitter
+from llama_index.core.response_synthesizers import CompactAndRefine
+from llama_index.core.schema import NodeWithScore
+from llama_index.postprocessor.rankgpt_rerank import RankGPTRerank
+from workflows import Context, Workflow, step
+from workflows.events import Event, StartEvent, StopEvent
+from llama_index.core import (VectorStoreIndex, SimpleDirectoryReader, StorageContext, load_index_from_storage)
+from llama_index.core import Settings 
+from healthagent.utils.logger import setup_logger
+from healthagent.utils.setup import DATA_DIR, STORAGE_DIR
+
+
+logger = setup_logger()
+
+
+class RetrieverEvent(Event):
+    """Result of running retrieval"""
+
+    nodes: list[NodeWithScore]
+
+
+class RerankEvent(Event):
+    """Result of running reranking on retrieved nodes"""
+
+    nodes: list[NodeWithScore]
+
+ 
+class RAGWorkflow(Workflow):
+    def __init__(self, index: VectorStoreIndex, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.index = index
+
+    @step
+    async def retrieve(self, ctx: Context, ev: StartEvent) -> RetrieverEvent | None:
+        "Entry point for RAG, triggered by a StartEvent with `query`."
+        
+        logger.info(f"Retrieving nodes for query: {ev.get('query')}")
+        query = ev.get("query")
+        top_k = ev.get("top_k", 5)
+        top_n = ev.get("top_n", 3)
+
+        if not query:
+            raise ValueError("Query is required!")
+
+        
+        await ctx.store.set("query", query)
+        await ctx.store.set("top_k", top_k)
+        await ctx.store.set("top_n", top_n)
+        
+
+        retriever = self.index.as_retriever(similarity_top_k=top_k)
+        nodes = await retriever.aretrieve(query)
+        return RetrieverEvent(nodes=nodes)
+
+    @step
+    async def rerank(self, ctx: Context, ev: RetrieverEvent) -> RerankEvent:
+        top_n = await ctx.store.get("top_n")
+        query = await ctx.store.get("query")
+
+        ranker = RankGPTRerank(top_n=top_n, llm=Settings.llm)
+
+        try:
+            new_nodes = ranker.postprocess_nodes(ev.nodes, query_str=query)
+        except Exception:
+            new_nodes = ev.nodes
+        return RerankEvent(nodes=new_nodes)
+
+    @step
+    async def synthesize(self, ctx: Context, ev: RerankEvent) -> StopEvent:
+        """Return a response using reranked nodes."""
+         
+        synthesizer = CompactAndRefine(llm=Settings.llm)
+        query = await ctx.store.get("query", default=None)
+
+        response = await synthesizer.asynthesize(query, nodes=ev.nodes)
+
+        return StopEvent(result=str(response))
+
+        
+
+class Retrieval:
+
+    def __init__(self):
+        if not os.path.exists(STORAGE_DIR):
+            logger.info("Building index.")
+            self.index = self.build_index()
+        else:
+            logger.info("Loading index.")
+            self.index = self.load_index()
+
+
+    #TODO: Change to mass index scheme
+    def build_index(self):
+        documents = SimpleDirectoryReader(DATA_DIR).load_data()
+        index = VectorStoreIndex.from_documents(documents)
+        index.storage_context.persist(STORAGE_DIR)
+
+        return index
+
+
+    #TODO: Change to mass index scheme
+    def load_index(self):
+        storage_context = StorageContext.from_defaults(persist_dir=STORAGE_DIR)
+        return load_index_from_storage(storage_context)
+
+
+    def query_context(self, text: str) -> str:
+        query_engine = self.index.as_query_engine()
+        result = query_engine.query(text)
+        return str(result)
+
+    def workflow(self) -> RAGWorkflow:
+        return RAGWorkflow(index=self.index, timeout=120.0)

healthagent/slm/routing.py

@@ -0,0 +1,25 @@
+from langgraph.graph import StateGraph, MessagesState, START, END
+
+class RouteDecision:
+    def __init__(self):
+        self.graph = StateGraph(MessagesState)
+        self.graph.add_node("reverse", self._reversal)
+        self.graph.add_node("verify", self._verify)
+        self.graph.add_edge(START, "_reversal")
+        #self.graph.add_conditional_edges("reverse", "verify") #TODO: Improve on the conditional edges
+        self.graph.add_edge("_reversal", END)
+
+    def _compile(self):
+        return self.graph.compile()
+
+    def _reversal(self, state: MessagesState):
+        """Responsible for reversing the prediction decision"""
+        return state["reverse"] = "Prediction Decision"
+
+    def _verify(self, state: MessagesState):
+        """Responsible for verifying the prediction from reversal"""
+        return state["_reversal"] = True
+
+    #TODO: Read the docs nodes into routing graphs
+    def _nodes_graph(self):
+        pass

healthagent/tools/document_nodes.py

@@ -0,0 +1,17 @@
+from src.utils.setup import BREAKPOINTS
+from llama_index.core import SimpleDirectoryReader
+from llama_index.core.node_parser import SentenceSplitter
+from llama_index.core.storage.docstore import SimpleDocumentStore
+
+def bake_docs_nodes(file_stream):
+    documents = SimpleDirectoryReader(file_stream).load_data()
+    node_parser  = SentenceSplitter(chunk_size=512)
+    base_nodes = node_parser.get_nodes_from_documents(documents)
+    return base_nodes
+
+
+#TODO: Responsible for serializing and caching documents into node docs.
+if __name__ == "__main__":
+    docstore = SimpleDocumentStore()
+    nodes = bake_docs_nodes(BREAKPOINTS)
+    docstore.add_documents(nodes)

healthagent/tools/explainer.py

@@ -0,0 +1,51 @@
+import pandas as pd
+import numpy as np
+import joblib
+from lime.lime_tabular import LimeTabularExplainer
+from healthagent.utils.setup import XIA_MODEL
+from healthagent.utils.data_lib import data_seralize
+from healthagent.utils.parser import parse_feature_contributions
+from healthagent.utils.logger import setup_logger
+
+logger = setup_logger()
+
+class XiaExplainer:
+  def __init__(self):
+    logger.info("Configuring explainer model.")
+    self.config = joblib.load(XIA_MODEL)
+    self.explainer = self.initialize_explainer()
+    logger.info("Explainer model configured sucessfully.")
+
+
+  def initialize_explainer(self):
+      return LimeTabularExplainer(
+          training_data=self.config["X_train"].values,
+          feature_names=self.config["feature_names"],
+          class_names=self.config["class_names"],
+          mode="classification",
+          verbose=False
+      )
+
+  def explain_prediction(self, model, data, parse=True, top_k=2):
+      feature_names=self.config["feature_names"]
+      class_names=self.config["class_names"]
+      sample_input= data_seralize(data)
+
+      probas = model.predict_proba(sample_input)[0]
+      idx = probas.argmax()
+      exp = self.explainer.explain_instance(
+          data_row= sample_input.values[0].astype(np.float16),
+          predict_fn=lambda x: model.predict_proba(pd.DataFrame(x, columns=feature_names)),
+          num_features=len(feature_names)
+      )
+      contributions = exp.as_list()
+      top = sorted(contributions, key=lambda x: abs(x[1]), reverse=True)[:top_k]
+
+      if parse:
+        return parse_feature_contributions(top)
+
+      return {
+          "predicted_class": class_names[idx],
+          "confidence": probas[idx],
+          "top_contributing_features": top
+      }

healthagent/types.py

@@ -0,0 +1,11 @@
+from pydantic import BaseModel, Field
+
+class ThrusterData(BaseModel):
+    """Schema for prediction input parameters."""
+    driver_temperature: float = Field(description="Driver temperature in Celsius")
+    temperature: float = Field(description="Temperature in Celsius")
+    speed_rpm: float = Field(description="Speed in RPM")
+    current: float = Field(description="Current in Amperes")
+    voltage: float = Field(description="Voltage in Volts")
+    power: float = Field(description="Power in Watts")
+    thruster_id_encoded: float = Field(default=1.0, description="Encoded thruster ID (default=1.0)")

healthagent/utils/data_lib.py

@@ -0,0 +1,27 @@
+import pandas as pd
+import random
+
+def data_seralize(thruster_data):
+    return pd.DataFrame(data={
+         "voltage": [thruster_data['voltage']],
+          "current": [thruster_data['current']],
+          "power": [thruster_data['power']],
+          "temperature": [thruster_data['temperature']],
+          "driver_temperature": [thruster_data['driver_temperature']],
+          "speed_rpm": [thruster_data['speed_rpm']],
+          "thruster_id_encoded": [thruster_data['thruster_id_encoded']]
+    })
+
+
+# For testing purpose
+def fake_data():
+    data = {
+        "voltage": round(random.uniform(0.05, 0.15), 3),
+        "current": round(random.uniform(5.0, 8.0), 3),
+        "power": round(random.uniform(0.5, 1.0), 3),
+        "temperature": round(random.uniform(25.0, 40.0), 3),
+        "driver_temperature": round(random.uniform(28.0, 45.0), 3),
+        "speed_rpm": round(random.uniform(3500, 5000), 3),
+        "thruster_id_encoded": float(random.randint(1, 8))
+    }
+    return data

healthagent/utils/logger.py

@@ -0,0 +1,13 @@
+import logging
+
+def setup_logger(name="Custom Logging System") -> logging.Logger:
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.INFO)
+
+    if not logger.handlers:
+        ch = logging.StreamHandler()
+        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
+        ch.setFormatter(formatter)
+        logger.addHandler(ch)
+
+    return logger

healthagent/utils/parser.py

@@ -0,0 +1,20 @@
+import re
+
+def parse_feature_contributions(contributions, to_text=True):
+    parsed = []
+    for label, value in contributions:
+        match = re.search(r"<\s*(\w+)\s*<=", label)
+        if match:
+            feature = match.group(1)
+        else:
+            parts = re.match(r"(.+?)\s*(<=|>=|<|>|=)\s*(.+)", label)
+            feature = parts.group(1).strip() if parts else label.strip()
+        parsed.append({
+            "feature": feature,
+            "contribution": f"{'+' if value >= 0 else '-'}{abs(value):.3f}"
+        })
+    if to_text:
+        response = [item['feature'] for index, item in enumerate(parsed)]
+        return ", ".join(response)
+
+    return parsed

healthagent/utils/setup.py

@@ -0,0 +1,18 @@
+import os
+import yaml
+from dotenv import load_dotenv
+
+load_dotenv()
+
+with open("healthagent/config/model_config.yaml", "r") as file:
+    config = yaml.safe_load(file)
+
+
+CHUNK_SIZE = config["chunk"]["size"]
+CHUNK_OVERLAP = config["chunk"]["overlap"]
+MODEL_NAME = config["model"]["language_name"]
+EMBED_MODEL = config["model"]["embedding"]
+PREDICTIVE_MODEL = config["model"]["predictive_model"]
+XIA_MODEL = config["model"]["xia_model"]
+DATA_DIR = config["documents"]["breakpoints"]
+STORAGE_DIR = config["storage"]

pyproject.toml

@@ -0,0 +1,48 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "healthagent"
+version = "0.1.0"
+description = "AI Health Agent for the Recon Robot"
+authors = [
+    { name = "Henry Asiedu", email = "henry.asiedu@fortressaisolutions.com" },
+    { name = "Amanda Ofori", email = "amanda.ofori@fortressaisolutions.com" }
+]
+readme = "README.md"
+requires-python = ">=3.9"
+maintainers = [
+    { name = "Henry Asiedu", email = "henry.asiedu@fortressaisolutions.com" },
+    { name = "Amanda Ofori", email = "amanda.ofori@fortressaisolutions.com" }
+]
+
+dependencies = [
+    "numpy",
+    "pandas",
+    "matplotlib",
+    "scikit-learn",
+    "seaborn",
+    "llama-index-core==0.13.3",
+    "llama-index",
+    "llama-index-llms-ollama",
+    "llama-index-embeddings-ollama",
+    "lime",
+    "pyyaml",
+    "joblib",
+    "fastapi",
+    "uvicorn",
+    "pydantic",
+    "llama-index-postprocessor-rankgpt-rerank>=0.2.0",
+    "llama-index-llms-mistralai==0.7.1",
+    "llama-index-embeddings-mistralai"
+]
+
+[project.optional-dependencies]
+dev = [
+    "streamlit"
+]
+
+[tool.setuptools]
+
+packages = ["healthagent"]