dlaima commited on
Commit
95f20ba
·
verified ·
1 Parent(s): 39211e6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +10 -7
app.py CHANGED
@@ -52,18 +52,21 @@ class MyAgent(CodeAgent):
52
  def __init__(self):
53
  model = HfApiModel(
54
  model="cognitivecomputations/dolphin-2.6-mixtral-8x7b",
55
- api_key=os.getenv("HF_API_TOKEN", "").strip(),
56
  )
57
  super().__init__(model=model, tools=tools)
58
  self.system_prompt = SYSTEM_PROMPT
59
 
60
  def __call__(self, task):
61
- # Inject the system prompt manually
62
- if isinstance(task, dict) and "question" in task:
63
- task["messages"] = [
64
- {"role": "system", "content": self.system_prompt},
65
- {"role": "user", "content": task["question"]}
66
- ]
 
 
 
67
  return super().__call__(task)
68
 
69
 
 
52
  def __init__(self):
53
  model = HfApiModel(
54
  model="cognitivecomputations/dolphin-2.6-mixtral-8x7b",
55
+ api_key=os.getenv("HF_API_TOKEN", "").strip()
56
  )
57
  super().__init__(model=model, tools=tools)
58
  self.system_prompt = SYSTEM_PROMPT
59
 
60
  def __call__(self, task):
61
+ # If task contains messages, prepend system prompt
62
+ if isinstance(task, dict):
63
+ if "question" in task:
64
+ task["messages"] = [
65
+ {"role": "system", "content": self.system_prompt},
66
+ {"role": "user", "content": task["question"]}
67
+ ]
68
+ elif "messages" in task:
69
+ task["messages"].insert(0, {"role": "system", "content": self.system_prompt})
70
  return super().__call__(task)
71
 
72