drrobot9 commited on
Commit
ddd60c5
·
verified ·
1 Parent(s): 85ce539

push updated backend changes and auto start buiding

Browse files
Files changed (1) hide show
  1. agents.py +6 -19
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)) # Exists but unused
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
- # Format with FormattingAgent
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)