support local llm server

README.md

@@ -78,6 +78,11 @@ To run Agent Laboratory in copilot mode, simply set the copilot-mode flag to `"t
 
 `python ai_lab_repo.py --api-key "API_KEY_HERE" --llm-backend "o1-mini" --research-topic "YOUR RESEARCH IDEA" --copilot-mode "true"`
 
+### Run with local LLM server
+
+`python ai_lab_repo.py --api-key "ollama" --llm-backend "qwen2.5-coder" --research-topic "YOUR RESEARCH IDEA" --copilot-mode "true" --base-url http://localhost:11434/v1`
+or 
+`python ai_lab_repo.py --api-key "llama.cpp" --llm-backend "qwen2.5-coder" --research-topic "YOUR RESEARCH IDEA" --copilot-mode "true" --base-url http://localhost:8080/v1`
 -----
 ## Tips for better research outcomes
 

agents.py

@@ -32,7 +32,7 @@ def extract_json_between_markers(llm_output):
 
 
 
-def get_score(outlined_plan, latex, reward_model_llm, reviewer_type=None, attempts=3, openai_api_key=None):
+def get_score(outlined_plan, latex, reward_model_llm, reviewer_type=None, attempts=3, openai_api_key=None, base_url=''):
     e = str()
     for _attempt in range(attempts):
         try:
@@ -144,7 +144,7 @@ def get_score(outlined_plan, latex, reward_model_llm, reviewer_type=None, attemp
                 openai_api_key=openai_api_key,
                 prompt=(
                     f"Outlined in the following text is the research plan that the machine learning engineer was tasked with building: {outlined_plan}\n\n"
-                    f"The following text is the research latex that the model produced: \n{latex}\n\n"), temp=0.0)
+                    f"The following text is the research latex that the model produced: \n{latex}\n\n"), temp=0.0, base_url=base_url)
             review_json = extract_json_between_markers(scoring)
 
             overall = int(review_json["Overall"]) / 10
@@ -181,27 +181,28 @@ def get_score(outlined_plan, latex, reward_model_llm, reviewer_type=None, attemp
 
 
 class ReviewersAgent:
-    def __init__(self, model="gpt-4o-mini", notes=None, openai_api_key=None):
+    def __init__(self, model="gpt-4o-mini", notes=None, openai_api_key=None, base_url=''):
         if notes is None: self.notes = []
         else: self.notes = notes
         self.model = model
         self.openai_api_key = openai_api_key
+        self.base_url=base_url
 
     def inference(self, plan, report):
         reviewer_1 = "You are a harsh but fair reviewer and expect good experiments that lead to insights for the research topic."
-        review_1 = get_score(outlined_plan=plan, latex=report, reward_model_llm=self.model, reviewer_type=reviewer_1, openai_api_key=self.openai_api_key)
+        review_1 = get_score(outlined_plan=plan, latex=report, reward_model_llm=self.model, reviewer_type=reviewer_1, openai_api_key=self.openai_api_key, base_url=self.base)
 
         reviewer_2 = "You are a harsh and critical but fair fair reviewer who is looking for idea that would be impactful in the field."
-        review_2 = get_score(outlined_plan=plan, latex=report, reward_model_llm=self.model, reviewer_type=reviewer_2, openai_api_key=self.openai_api_key)
+        review_2 = get_score(outlined_plan=plan, latex=report, reward_model_llm=self.model, reviewer_type=reviewer_2, openai_api_key=self.openai_api_key, base_url=self.base)
 
         reviewer_3 = "You are a harsh but fair open-minded reviewer that is looking for novel ideas that have not been proposed before."
-        review_3 = get_score(outlined_plan=plan, latex=report, reward_model_llm=self.model, reviewer_type=reviewer_3, openai_api_key=self.openai_api_key)
+        review_3 = get_score(outlined_plan=plan, latex=report, reward_model_llm=self.model, reviewer_type=reviewer_3, openai_api_key=self.openai_api_key, base_url=self.base)
 
         return f"Reviewer #1:\n{review_1}, \nReviewer #2:\n{review_2}, \nReviewer #3:\n{review_3}"
 
 
 class BaseAgent:
-    def __init__(self, model="gpt-4o-mini", notes=None, max_steps=100, openai_api_key=None):
+    def __init__(self, model="gpt-4o-mini", notes=None, max_steps=100, openai_api_key=None, base_url=''):
         if notes is None: self.notes = []
         else: self.notes = notes
         self.max_steps = max_steps
@@ -222,6 +223,7 @@ def __init__(self, model="gpt-4o-mini", notes=None, max_steps=100, openai_api_ke
         self.prev_results_code = str()
         self.prev_interpretation = str()
         self.openai_api_key = openai_api_key
+        self.base_url=base_url
 
         self.second_round = False
         self.max_hist_len = 15
@@ -251,7 +253,7 @@ def inference(self, research_topic, phase, step, feedback="", temp=None):
             f"Current Step #{step}, Phase: {phase}\n{complete_str}\n"
             f"[Objective] Your goal is to perform research on the following topic: {research_topic}\n"
             f"Feedback: {feedback}\nNotes: {notes_str}\nYour previous command was: {self.prev_comm}. Make sure your new output is very different.\nPlease produce a single command below:\n")
-        model_resp = query_model(model_str=self.model, system_prompt=sys_prompt, prompt=prompt, temp=temp, openai_api_key=self.openai_api_key)
+        model_resp = query_model(model_str=self.model, system_prompt=sys_prompt, prompt=prompt, temp=temp, openai_api_key=self.openai_api_key, base_url=self.base_url)
         print("^"*50, phase, "^"*50)
         model_resp = self.clean_text(model_resp)
         self.prev_comm = model_resp
@@ -291,8 +293,8 @@ def example_command(self, phase):
 
 
 class ProfessorAgent(BaseAgent):
-    def __init__(self, model="gpt4omini", notes=None, max_steps=100, openai_api_key=None):
-        super().__init__(model, notes, max_steps, openai_api_key)
+    def __init__(self, model="gpt4omini", notes=None, max_steps=100, openai_api_key=None, base_url=''):
+        super().__init__(model, notes, max_steps, openai_api_key, base_url)
         self.phases = ["report writing"]
 
     def generate_readme(self):
@@ -301,7 +303,7 @@ def generate_readme(self):
         prompt = (
             f"""History: {history_str}\n{'~' * 10}\n"""
             f"Please produce the readme below in markdown:\n")
-        model_resp = query_model(model_str=self.model, system_prompt=sys_prompt, prompt=prompt, openai_api_key=self.openai_api_key)
+        model_resp = query_model(model_str=self.model, system_prompt=sys_prompt, prompt=prompt, openai_api_key=self.openai_api_key, base_url=self.base_url)
         return model_resp.replace("```markdown", "")
 
     def context(self, phase):
@@ -357,8 +359,8 @@ def role_description(self):
 
 
 class PostdocAgent(BaseAgent):
-    def __init__(self, model="gpt4omini", notes=None, max_steps=100, openai_api_key=None):
-        super().__init__(model, notes, max_steps, openai_api_key)
+    def __init__(self, model="gpt4omini", notes=None, max_steps=100, openai_api_key=None, base_url=''):
+        super().__init__(model, notes, max_steps, openai_api_key, base_url)
         self.phases = ["plan formulation", "results interpretation"]
 
     def context(self, phase):
@@ -433,8 +435,8 @@ def role_description(self):
 
 
 class MLEngineerAgent(BaseAgent):
-    def __init__(self, model="gpt4omini", notes=None, max_steps=100, openai_api_key=None):
-        super().__init__(model, notes, max_steps, openai_api_key)
+    def __init__(self, model="gpt4omini", notes=None, max_steps=100, openai_api_key=None, base_url=''):
+        super().__init__(model, notes, max_steps, openai_api_key, base_url)
         self.phases = [
             "data preparation",
             "running experiments",
@@ -498,8 +500,8 @@ def role_description(self):
 
 
 class PhDStudentAgent(BaseAgent):
-    def __init__(self, model="gpt4omini", notes=None, max_steps=100, openai_api_key=None):
-        super().__init__(model, notes, max_steps, openai_api_key)
+    def __init__(self, model="gpt4omini", notes=None, max_steps=100, openai_api_key=None, base_url=''):
+        super().__init__(model, notes, max_steps, openai_api_key,base_url)
         self.phases = [
             "literature review",
             "plan formulation",
@@ -579,7 +581,7 @@ def requirements_txt(self):
         prompt = (
             f"""History: {history_str}\n{'~' * 10}\n"""
             f"Please produce the requirements.txt below in markdown:\n")
-        model_resp = query_model(model_str=self.model, system_prompt=sys_prompt, prompt=prompt, openai_api_key=self.openai_api_key)
+        model_resp = query_model(model_str=self.model, system_prompt=sys_prompt, prompt=prompt, openai_api_key=self.openai_api_key, base_url=self.base_url)
         return model_resp
 
     def example_command(self, phase):

ai_lab_repo.py

@@ -11,7 +11,8 @@
 
 
 class LaboratoryWorkflow:
-    def __init__(self, research_topic, openai_api_key, max_steps=100, num_papers_lit_review=5, agent_model_backbone=f"{DEFAULT_LLM_BACKBONE}", notes=list(), human_in_loop_flag=None, compile_pdf=True, mlesolver_max_steps=3, papersolver_max_steps=5):
+    def __init__(self, research_topic, openai_api_key, max_steps=100, num_papers_lit_review=5, agent_model_backbone=f"{DEFAULT_LLM_BACKBONE}",
+                 notes=list(), human_in_loop_flag=None, compile_pdf=True, mlesolver_max_steps=3, papersolver_max_steps=5, base_url=''):
         """
         Initialize laboratory workflow
         @param research_topic: (str) description of research idea to explore
@@ -25,6 +26,7 @@ def __init__(self, research_topic, openai_api_key, max_steps=100, num_papers_lit
         self.max_steps = max_steps
         self.compile_pdf = compile_pdf
         self.openai_api_key = openai_api_key
+        self.base_url=base_url
         self.research_topic = research_topic
         self.model_backbone = agent_model_backbone
         self.num_papers_lit_review = num_papers_lit_review
@@ -79,11 +81,11 @@ def __init__(self, research_topic, openai_api_key, max_steps=100, num_papers_lit
 
         self.save = True
         self.verbose = True
-        self.reviewers = ReviewersAgent(model=self.model_backbone, notes=self.notes, openai_api_key=self.openai_api_key)
-        self.phd = PhDStudentAgent(model=self.model_backbone, notes=self.notes, max_steps=self.max_steps, openai_api_key=self.openai_api_key)
-        self.postdoc = PostdocAgent(model=self.model_backbone, notes=self.notes, max_steps=self.max_steps, openai_api_key=self.openai_api_key)
-        self.professor = ProfessorAgent(model=self.model_backbone, notes=self.notes, max_steps=self.max_steps, openai_api_key=self.openai_api_key)
-        self.ml_engineer = MLEngineerAgent(model=self.model_backbone, notes=self.notes, max_steps=self.max_steps, openai_api_key=self.openai_api_key)
+        self.reviewers = ReviewersAgent(model=self.model_backbone, notes=self.notes, openai_api_key=self.openai_api_key, base_url=self.base_url)
+        self.phd = PhDStudentAgent(model=self.model_backbone, notes=self.notes, max_steps=self.max_steps, openai_api_key=self.openai_api_key, base_url=self.base_url)
+        self.postdoc = PostdocAgent(model=self.model_backbone, notes=self.notes, max_steps=self.max_steps, openai_api_key=self.openai_api_key, base_url=self.base_url)
+        self.professor = ProfessorAgent(model=self.model_backbone, notes=self.notes, max_steps=self.max_steps, openai_api_key=self.openai_api_key, base_url=self.base_url)
+        self.ml_engineer = MLEngineerAgent(model=self.model_backbone, notes=self.notes, max_steps=self.max_steps, openai_api_key=self.openai_api_key, base_url=self.base_url)
 
         # remove previous files
         remove_figures()
@@ -240,7 +242,8 @@ def report_writing(self):
         # instantiate mle-solver
         from papersolver import PaperSolver
         self.reference_papers = []
-        solver = PaperSolver(notes=report_notes, max_steps=self.papersolver_max_steps, plan=lab.phd.plan, exp_code=lab.phd.results_code, exp_results=lab.phd.exp_results, insights=lab.phd.interpretation, lit_review=lab.phd.lit_review, ref_papers=self.reference_papers, topic=research_topic, openai_api_key=self.openai_api_key, llm_str=self.model_backbone["report writing"], compile_pdf=compile_pdf)
+        solver = PaperSolver(notes=report_notes, max_steps=self.papersolver_max_steps, plan=lab.phd.plan, exp_code=lab.phd.results_code, exp_results=lab.phd.exp_results, insights=lab.phd.interpretation, lit_review=lab.phd.lit_review, ref_papers=self.reference_papers, topic=research_topic, openai_api_key=self.openai_api_key,
+                             llm_str=self.model_backbone["report writing"], compile_pdf=compile_pdf, base_url=self.base_url)
         # run initialization for solver
         solver.initial_solve()
         # run solver for N mle optimization steps
@@ -603,17 +606,24 @@ def parse_arguments():
         help='Total number of paper-solver steps'
     )
 
+    parser.add_argument(
+        '--base-url',
+        type=str,
+        default="",
+        help='Set to different url if you are using a custom server, e.g. http://localhost:11434 if using ollama, or http://localhost:8080/v1 if using llama.cpp.'
+    )
+
 
     return parser.parse_args()
 
 
 if __name__ == "__main__":
     args = parse_arguments()
-
     llm_backend = args.llm_backend
     human_mode = args.copilot_mode.lower() == "true"
     compile_pdf = args.compile_latex.lower() == "true"
     load_existing = args.load_existing.lower() == "true"
+    base_url=args.base_url
     try:
         num_papers_lit_review = int(args.num_papers_lit_review.lower())
     except Exception:
@@ -711,6 +721,7 @@ def parse_arguments():
             num_papers_lit_review=num_papers_lit_review,
             papersolver_max_steps=papersolver_max_steps,
             mlesolver_max_steps=mlesolver_max_steps,
+            base_url=base_url
         )
 
     lab.perform_research()

inference.py

@@ -8,6 +8,7 @@
 
 encoding = tiktoken.get_encoding("cl100k_base")
 
+
 def curr_cost_est():
     costmap_in = {
         "gpt-4o": 2.50 / 1000000,
@@ -17,15 +18,18 @@ def curr_cost_est():
         "claude-3-5-sonnet": 3.00 / 1000000,
     }
     costmap_out = {
-        "gpt-4o": 10.00/ 1000000,
+        "gpt-4o": 10.00 / 1000000,
         "gpt-4o-mini": 0.6 / 1000000,
         "o1-preview": 60.00 / 1000000,
         "o1-mini": 12.00 / 1000000,
         "claude-3-5-sonnet": 12.00 / 1000000,
     }
-    return sum([costmap_in[_]*TOKENS_IN[_] for _ in TOKENS_IN]) + sum([costmap_out[_]*TOKENS_OUT[_] for _ in TOKENS_OUT])
+    return sum([costmap_in[_] * TOKENS_IN[_] for _ in TOKENS_IN]) + sum(
+        [costmap_out[_] * TOKENS_OUT[_] for _ in TOKENS_OUT])
+
 
-def query_model(model_str, prompt, system_prompt, openai_api_key=None, anthropic_api_key=None, tries=5, timeout=5.0, temp=None, print_cost=True, version="1.5"):
+def query_model(model_str, prompt, system_prompt, openai_api_key=None, anthropic_api_key=None, tries=5, timeout=5.0,
+                temp=None, print_cost=True, version="1.5", base_url=''):
     preloaded_api = os.getenv('OPENAI_API_KEY')
     if openai_api_key is None and preloaded_api is not None:
         openai_api_key = preloaded_api
@@ -36,6 +40,7 @@ def query_model(model_str, prompt, system_prompt, openai_api_key=None, anthropic
         os.environ["OPENAI_API_KEY"] = openai_api_key
     if anthropic_api_key is not None:
         os.environ["ANTHROPIC_API_KEY"] = anthropic_api_key
+    encoding=None
     for _ in range(tries):
         try:
             if model_str == "gpt-4o-mini" or model_str == "gpt4omini" or model_str == "gpt-4omini" or model_str == "gpt4o-mini":
@@ -124,23 +129,41 @@ def query_model(model_str, prompt, system_prompt, openai_api_key=None, anthropic
                     completion = client.chat.completions.create(
                         model="o1-preview", messages=messages)
                 answer = completion.choices[0].message.content
-
-            if model_str in ["o1-preview", "o1-mini", "claude-3.5-sonnet"]:
-                encoding = tiktoken.encoding_for_model("gpt-4o")
-            else: encoding = tiktoken.encoding_for_model(model_str)
+            else:
+                messages = [
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": prompt}]
+                client = OpenAI()
+                if base_url is not None and base_url != '':
+                    client.base_url = base_url
+                completion = client.chat.completions.create(
+                    model=f"{model_str}", messages=messages)
+                answer = completion.choices[0].message.content
+            try:
+                if model_str in ["o1-preview", "o1-mini", "claude-3.5-sonnet"]:
+                    encoding = tiktoken.encoding_for_model("gpt-4o")
+                else:
+                    encoding = tiktoken.encoding_for_model(model_str)
+            except Exception as e:
+                pass
+            if encoding is None:
+                # set a default encoding to by pass the error for now.
+                encoding = tiktoken.encoding_for_model('gpt-4')
             if model_str not in TOKENS_IN:
                 TOKENS_IN[model_str] = 0
                 TOKENS_OUT[model_str] = 0
             TOKENS_IN[model_str] += len(encoding.encode(system_prompt + prompt))
             TOKENS_OUT[model_str] += len(encoding.encode(answer))
             if print_cost:
-                print(f"Current experiment cost = ${curr_cost_est()}, ** Approximate values, may not reflect true cost")
+                try:
+                    print(f"Current experiment cost = ${curr_cost_est()}, ** Approximate values, may not reflect true cost")
+                except Exception as e:
+                    pass
             return answer
         except Exception as e:
             print("Inference Exception:", e)
             time.sleep(timeout)
             continue
     raise Exception("Max retries: timeout")
 
-
-#print(query_model(model_str="o1-mini", prompt="hi", system_prompt="hey"))
+# print(query_model(model_str="o1-mini", prompt="hi", system_prompt="hey"))