feat: ollama integration

README.md

@@ -1,6 +1,5 @@
 # Agent Laboratory: Using LLM Agents as Research Assistants
 
-
 <p align="center">
   <img src="media/AgentLabLogo.png" alt="Demonstration of the flow of AgentClinic" style="width: 99%;">
 </p>
@@ -15,7 +14,7 @@
 
 ## 📖 Overview
 
-- **Agent Laboratory** is an end-to-end autonomous research workflow meant to assist **you** as the human researcher toward **implementing your research ideas**. Agent Laboratory consists of specialized agents driven by large language models to support you through the entire research workflow—from conducting literature reviews and formulating plans to executing experiments and writing comprehensive reports. 
+- **Agent Laboratory** is an end-to-end autonomous research workflow meant to assist **you** as the human researcher toward **implementing your research ideas**. Agent Laboratory consists of specialized agents driven by large language models to support you through the entire research workflow—from conducting literature reviews and formulating plans to executing experiments and writing comprehensive reports.
 - This system is not designed to replace your creativity but to complement it, enabling you to focus on ideation and critical thinking while automating repetitive and time-intensive tasks like coding and documentation. By accommodating varying levels of computational resources and human involvement, Agent Laboratory aims to accelerate scientific discovery and optimize your research productivity.
 
 <p align="center">
@@ -30,48 +29,52 @@
   <img src="media/AgentLabWF.png" alt="Demonstration of the flow of AgentClinic" style="width: 99%;">
 </p>
 
-
 ### 👾 Currently supported models
 
-* **OpenAI**: o1, o1-preview, o1-mini, gpt-4o
-* **DeepSeek**: deepseek-chat (deepseek-v3)
+- **OpenAI**: o1, o1-preview, o1-mini, gpt-4o
+- **DeepSeek**: deepseek-chat (deepseek-v3)
 
 To select a specific llm set the flag `--llm-backend="llm_model"` for example `--llm-backend="gpt-4o"` or `--llm-backend="deepseek-chat"`. Please feel free to add a PR supporting new models according to your need!
 
 ## 🖥️ Installation
 
 ### Python venv option
 
-* We recommend using python 3.12
+- We recommend using python 3.12
 
 1. **Clone the GitHub Repository**: Begin by cloning the repository using the command:
+
 ```bash
 git clone git@github.com:SamuelSchmidgall/AgentLaboratory.git
 ```
 
 2. **Set up and Activate Python Environment**
+
 ```bash
 python -m venv venv_agent_lab
 ```
+
 - Now activate this environment:
+
 ```bash
 source venv_agent_lab/bin/activate
 ```
 
 3. **Install required libraries**
+
 ```bash
 pip install -r requirements.txt
 ```
 
 4. **Install pdflatex [OPTIONAL]**
+
 ```bash
 sudo apt install pdflatex
 ```
+
 - This enables latex source to be compiled by the agents.
 - **[IMPORTANT]** If this step cannot be run due to not having sudo access, pdf compiling can be turned off via running Agent Laboratory via setting the `--compile-latex` flag to false: `--compile-latex "false"`
 
-
-
 5. **Now run Agent Laboratory!**
 
 `python ai_lab_repo.py --api-key "API_KEY_HERE" --llm-backend "o1-mini" --research-topic "YOUR RESEARCH IDEA"`
@@ -86,18 +89,26 @@ 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"`
 
------
-## Tips for better research outcomes
+### Running with Ollama
 
+To run Agent Laboratory locally with ollama, use prefix `ollama:` with you local model.
+
+`python ai_lab_repo.py --llm-backend "ollama:model_name" --research-topic "YOUR RESEARCH IDEA"`
+
+Example:
+`python ai_lab_repo.py --llm-backend "ollama:deepseek-r1:1.5b" --research-topic "YOUR RESEARCH IDEA"`
+
+---
+
+## Tips for better research outcomes
 
 #### [Tip #1] 📝 Make sure to write extensive notes! 📝
 
 **Writing extensive notes is important** for helping your agent understand what you're looking to accomplish in your project, as well as any style preferences. Notes can include any experiments you want the agents to perform, providing API keys, certain plots or figures you want included, or anything you want the agent to know when performing research.
 
 This is also your opportunity to let the agent know **what compute resources it has access to**, e.g. GPUs (how many, what type of GPU, how many GBs), CPUs (how many cores, what type of CPUs), storage limitations, and hardware specs.
 
-In order to add notes, you must modify the task_notes_LLM structure inside of `ai_lab_repo.py`. Provided below is an example set of notes used for some of our experiments. 
-
+In order to add notes, you must modify the task_notes_LLM structure inside of `ai_lab_repo.py`. Provided below is an example set of notes used for some of our experiments.
 
 ```
 task_notes_LLM = [
@@ -124,7 +135,7 @@ task_notes_LLM = [
     ]
 ```
 
---------
+---
 
 #### [Tip #2] 🚀 Using more powerful models generally leads to better research 🚀
 
@@ -134,17 +145,15 @@ However, **it’s important to balance performance and cost-effectiveness**. Whi
 
 When resources are limited, **optimize by fine-tuning smaller models** on your specific dataset or combining pre-trained models with task-specific prompts to achieve the desired balance between performance and computational efficiency.
 
------
+---
 
 #### [Tip #3] ✅ You can load previous saves from checkpoints ✅
 
 **If you lose progress, internet connection, or if a subtask fails, you can always load from a previous state.** All of your progress is saved by default in the `state_saves` variable, which stores each individual checkpoint. Just pass the following arguments when running `ai_lab_repo.py`
 
 `python ai_lab_repo.py --api-key "API_KEY_HERE" --research-topic "YOUR RESEARCH IDEA" --llm-backend "o1-mini" --load-existing True --load-existing-path "state_saves/LOAD_PATH"`
 
------
-
-
+---
 
 #### [Tip #4] 🈯 If you are running in a language other than English 🈲
 
@@ -154,14 +163,12 @@ For example, if you are running in Chinese:
 
 `python ai_lab_repo.py --api-key "API_KEY_HERE" --research-topic "YOUR RESEARCH IDEA (in your language)" --llm-backend "o1-mini" --language "中文"`
 
-----
-
+---
 
 #### [Tip #5] 🌟 There is a lot of room for improvement 🌟
 
 There is a lot of room to improve this codebase, so if you end up making changes and want to help the community, please feel free to share the changes you've made! We hope this tool helps you!
 
-
 ## 📜 License
 
 Source Code Licensing: Our project's source code is licensed under the MIT License. This license permits the use, modification, and distribution of the code, subject to certain conditions outlined in the MIT License.
@@ -172,16 +179,14 @@ If you would like to get in touch, feel free to reach out to [sschmi46@jhu.edu](
 
 ## Reference / Bibtex
 
-
-
 ```bibtex
 @misc{schmidgall2025agentlaboratoryusingllm,
-      title={Agent Laboratory: Using LLM Agents as Research Assistants}, 
+      title={Agent Laboratory: Using LLM Agents as Research Assistants},
       author={Samuel Schmidgall and Yusheng Su and Ze Wang and Ximeng Sun and Jialian Wu and Xiaodong Yu and Jiang Liu and Zicheng Liu and Emad Barsoum},
       year={2025},
       eprint={2501.04227},
       archivePrefix={arXiv},
       primaryClass={cs.HC},
-      url={https://arxiv.org/abs/2501.04227}, 
+      url={https://arxiv.org/abs/2501.04227},
 }
 ```

ai_lab_repo.py

@@ -643,7 +643,7 @@ def parse_arguments():
     if args.deepseek_api_key is not None and os.getenv('DEEPSEEK_API_KEY') is None:
         os.environ["DEEPSEEK_API_KEY"] = args.deepseek_api_key
 
-    if not api_key and not deepseek_api_key:
+    if not api_key and not deepseek_api_key and not llm_backend.startswith('ollama:'):
         raise ValueError("API key must be provided via --api-key / -deepseek-api-key or the OPENAI_API_KEY / DEEPSEEK_API_KEY environment variable.")
 
     ##########################################################

inference.py

@@ -1,7 +1,7 @@
 import time, tiktoken
 from openai import OpenAI
 import openai
-import os, anthropic, json
+import os, anthropic, json, ollama
 
 TOKENS_IN = dict()
 TOKENS_OUT = dict()
@@ -29,17 +29,38 @@ def curr_cost_est():
     }
     return sum([costmap_in[_]*TOKENS_IN[_] for _ in TOKENS_IN]) + sum([costmap_out[_]*TOKENS_OUT[_] for _ in TOKENS_OUT])
 
+def compute_tokens(model_str, prompt, system_prompt, answer, print_cost): 
+    try:           
+        if model_str in ["o1-preview", "o1-mini", "claude-3.5-sonnet", "o1"]:
+            encoding = tiktoken.encoding_for_model("gpt-4o")
+        elif model_str in ["deepseek-chat"]:
+            encoding = tiktoken.encoding_for_model("cl100k_base")
+        else:
+            encoding = tiktoken.encoding_for_model(model_str)
+        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")
+    except Exception as e:
+                if print_cost:
+                    print(f"Cost approximation has an error? {e}")
+
 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"):
-    preloaded_api = os.getenv('OPENAI_API_KEY')
-    if openai_api_key is None and preloaded_api is not None:
-        openai_api_key = preloaded_api
-    if openai_api_key is None and anthropic_api_key is None:
-        raise Exception("No API key provided in query_model function")
-    if openai_api_key is not None:
-        openai.api_key = openai_api_key
-        os.environ["OPENAI_API_KEY"] = openai_api_key
-    if anthropic_api_key is not None:
-        os.environ["ANTHROPIC_API_KEY"] = anthropic_api_key
+
+    if not model_str.startswith('ollama:'): 
+        preloaded_api = os.getenv('OPENAI_API_KEY')
+        if openai_api_key is None and preloaded_api is not None:
+            openai_api_key = preloaded_api
+        if openai_api_key is None and anthropic_api_key is None:
+            raise Exception("No API key provided in query_model function")
+        if openai_api_key is not None:
+            openai.api_key = openai_api_key
+            os.environ["OPENAI_API_KEY"] = openai_api_key
+        if anthropic_api_key is not None:
+            os.environ["ANTHROPIC_API_KEY"] = anthropic_api_key
     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":
@@ -159,24 +180,22 @@ 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
-
-            try:
-                if model_str in ["o1-preview", "o1-mini", "claude-3.5-sonnet", "o1"]:
-                    encoding = tiktoken.encoding_for_model("gpt-4o")
-                elif model_str in ["deepseek-chat"]:
-                    encoding = tiktoken.encoding_for_model("cl100k_base")
-                else:
-                    encoding = tiktoken.encoding_for_model(model_str)
-                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")
-            except Exception as e:
-                if print_cost:
-                    print(f"Cost approximation has an error? {e}")
+            elif model_str.startswith("ollama:"): 
+
+                response: ollama.ChatResponse = ollama.chat(
+                model=model_str[7:], 
+                messages=[ 
+                    { "role": "system", "content": system_prompt }, 
+                    { "role": "user", "content":  prompt }
+                ]
+                )
+
+                answer = response.message.content 
+
+            # skip token computation for models run via ollama
+            if not model_str.startswith("ollama:"): 
+                compute_tokens(model_str,prompt, system_prompt, answer, print_cost)
+
             return answer
         except Exception as e:
             print("Inference Exception:", e)
@@ -185,4 +204,5 @@ def query_model(model_str, prompt, system_prompt, openai_api_key=None, anthropic
     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"))
+# print(query_model(model_str="ollama:deepseek-r1:1.5b", prompt="hi", system_prompt="hey"))

requirements.txt

@@ -68,6 +68,7 @@ nest-asyncio==1.6.0
 networkx==3.2.1
 nltk==3.9.1
 numpy==2.0.2
+ollama==0.4.7
 openai==1.55.1
 opt_einsum==3.4.0
 optree==0.13.1