push updated backend changes and auto start buiding
Browse files
agents.py
CHANGED
|
@@ -1,3 +1,4 @@
|
|
|
|
|
| 1 |
import torch
|
| 2 |
from transformers import AutoModelForCausalLM, AutoTokenizer
|
| 3 |
from autogen import AssistantAgent, UserProxyAgent
|
|
@@ -10,8 +11,8 @@ BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
| 10 |
if BASE_DIR not in sys.path:
|
| 11 |
sys.path.insert(0, BASE_DIR)
|
| 12 |
|
| 13 |
-
# Load BioMistral once
|
| 14 |
|
|
|
|
| 15 |
class BioMistralModel:
|
| 16 |
def __init__(self, model_name=LLM_MODEL, device=None):
|
| 17 |
print(f"[BioMistralModel] Loading model: {model_name}")
|
|
@@ -38,56 +39,42 @@ class BioMistralModel:
|
|
| 38 |
)
|
| 39 |
|
| 40 |
text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
|
| 41 |
-
|
| 42 |
return text.split("Answer:", 1)[-1].strip()
|
| 43 |
|
| 44 |
|
| 45 |
-
|
| 46 |
# Formatting Agent
|
| 47 |
-
|
| 48 |
class FormattingAgent(AssistantAgent):
|
| 49 |
def __init__(self, name="FormattingAgent", **kwargs):
|
| 50 |
super().__init__(name=name, **kwargs)
|
| 51 |
|
| 52 |
def format_text(self, text: str) -> str:
|
| 53 |
-
|
| 54 |
cleaned = " ".join(text.split())
|
| 55 |
-
|
| 56 |
if cleaned:
|
| 57 |
cleaned = cleaned[0].upper() + cleaned[1:]
|
| 58 |
return cleaned
|
| 59 |
|
| 60 |
|
| 61 |
-
|
| 62 |
# Tutor Agent
|
| 63 |
-
|
| 64 |
class TutorAgent(AssistantAgent):
|
| 65 |
def __init__(self, name="TutorAgent", **kwargs):
|
| 66 |
super().__init__(name=name, **kwargs)
|
| 67 |
self.model = BioMistralModel()
|
| 68 |
self.format_agent = FormattingAgent()
|
| 69 |
-
self.rag_agent = RAGAgent(vectorstore_dir=str(VECTORSTORE_DIR)) #
|
| 70 |
|
| 71 |
def process_query(self, query: str) -> str:
|
| 72 |
print(f"[TutorAgent] Received query: {query}")
|
| 73 |
|
| 74 |
-
# Direct BioMistral answer
|
| 75 |
answer = self.model.generate_answer(query)
|
| 76 |
confidence = self.estimate_confidence(answer)
|
| 77 |
|
| 78 |
print(f"[TutorAgent] Confidence: {confidence:.2f}")
|
| 79 |
if confidence < CONFIDENCE_THRESHOLD:
|
| 80 |
print("[TutorAgent] Confidence low, but still using BioMistral (RAG unused).")
|
| 81 |
-
|
| 82 |
|
| 83 |
-
|
| 84 |
-
final_answer = self.format_agent.format_text(answer)
|
| 85 |
-
return final_answer
|
| 86 |
|
| 87 |
def estimate_confidence(self, answer: str) -> float:
|
| 88 |
-
"""
|
| 89 |
-
Dummy confidence estimator — could be replaced with actual logic.
|
| 90 |
-
"""
|
| 91 |
length = len(answer.strip())
|
| 92 |
if length > 100:
|
| 93 |
return 0.9
|
|
@@ -97,9 +84,9 @@ class TutorAgent(AssistantAgent):
|
|
| 97 |
return 0.5
|
| 98 |
|
| 99 |
|
| 100 |
-
#
|
| 101 |
# User Agent
|
| 102 |
-
|
| 103 |
class BioUser(UserProxyAgent):
|
| 104 |
def __init__(self, name="BioUser", **kwargs):
|
|
|
|
|
|
|
| 105 |
super().__init__(name=name, **kwargs)
|
|
|
|
| 1 |
+
# bioinformatics_ai/agents.py
|
| 2 |
import torch
|
| 3 |
from transformers import AutoModelForCausalLM, AutoTokenizer
|
| 4 |
from autogen import AssistantAgent, UserProxyAgent
|
|
|
|
| 11 |
if BASE_DIR not in sys.path:
|
| 12 |
sys.path.insert(0, BASE_DIR)
|
| 13 |
|
|
|
|
| 14 |
|
| 15 |
+
# Load BioMistral once
|
| 16 |
class BioMistralModel:
|
| 17 |
def __init__(self, model_name=LLM_MODEL, device=None):
|
| 18 |
print(f"[BioMistralModel] Loading model: {model_name}")
|
|
|
|
| 39 |
)
|
| 40 |
|
| 41 |
text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
|
|
|
|
| 42 |
return text.split("Answer:", 1)[-1].strip()
|
| 43 |
|
| 44 |
|
|
|
|
| 45 |
# Formatting Agent
|
|
|
|
| 46 |
class FormattingAgent(AssistantAgent):
|
| 47 |
def __init__(self, name="FormattingAgent", **kwargs):
|
| 48 |
super().__init__(name=name, **kwargs)
|
| 49 |
|
| 50 |
def format_text(self, text: str) -> str:
|
|
|
|
| 51 |
cleaned = " ".join(text.split())
|
|
|
|
| 52 |
if cleaned:
|
| 53 |
cleaned = cleaned[0].upper() + cleaned[1:]
|
| 54 |
return cleaned
|
| 55 |
|
| 56 |
|
|
|
|
| 57 |
# Tutor Agent
|
|
|
|
| 58 |
class TutorAgent(AssistantAgent):
|
| 59 |
def __init__(self, name="TutorAgent", **kwargs):
|
| 60 |
super().__init__(name=name, **kwargs)
|
| 61 |
self.model = BioMistralModel()
|
| 62 |
self.format_agent = FormattingAgent()
|
| 63 |
+
self.rag_agent = RAGAgent(vectorstore_dir=str(VECTORSTORE_DIR)) # safe conversion
|
| 64 |
|
| 65 |
def process_query(self, query: str) -> str:
|
| 66 |
print(f"[TutorAgent] Received query: {query}")
|
| 67 |
|
|
|
|
| 68 |
answer = self.model.generate_answer(query)
|
| 69 |
confidence = self.estimate_confidence(answer)
|
| 70 |
|
| 71 |
print(f"[TutorAgent] Confidence: {confidence:.2f}")
|
| 72 |
if confidence < CONFIDENCE_THRESHOLD:
|
| 73 |
print("[TutorAgent] Confidence low, but still using BioMistral (RAG unused).")
|
|
|
|
| 74 |
|
| 75 |
+
return self.format_agent.format_text(answer)
|
|
|
|
|
|
|
| 76 |
|
| 77 |
def estimate_confidence(self, answer: str) -> float:
|
|
|
|
|
|
|
|
|
|
| 78 |
length = len(answer.strip())
|
| 79 |
if length > 100:
|
| 80 |
return 0.9
|
|
|
|
| 84 |
return 0.5
|
| 85 |
|
| 86 |
|
|
|
|
| 87 |
# User Agent
|
|
|
|
| 88 |
class BioUser(UserProxyAgent):
|
| 89 |
def __init__(self, name="BioUser", **kwargs):
|
| 90 |
+
|
| 91 |
+
kwargs.setdefault("code_execution_config", {"use_docker": False})
|
| 92 |
super().__init__(name=name, **kwargs)
|