# DrSai 
A framework for rapid development of intelligent agent and multi-agent collaborative systems developed by the Dr.Sai team at the Institute of High Energy Physics.

<div align="center">
  <p>
      <img width="80%" src="assets/drsai.png" alt="Adaptive Logic Diagram">
  </p>
</div>

## 1. Features

- 1. Flexible switching of agent base models based on the [HepAI Platform](https://ai.ihep.ac.cn/).
- 2. Designed perception, thinking, memory, execution, and other behavioral functions for agents, with a plugin architecture for flexible expansion to meet various application scenarios.
- 3. Provides multiple multi-agent collaboration logic such as agent selection and hand-raising systems.
- 4. Provides a compatible OpenAI Chat and OpenAI ASSISTANTS standard backend interface for agent and multi-agent collaboration system interaction, allowing seamless integration with OpenAI-compatible frontend outputs. This allows the agent and multi-agent collaboration system to be deployed as a model or agent service.

## 2. Getting Started

### 2.1. Install the DrSai Multi-Agent Collaboration System

#### Install via pip

```shell
conda create -n drsai python==3.10
conda activate drsai
pip install DrSai -U
```

Configure API access keys and other environment variables for DDF2 platforms like [HepAI](https://ai.ihep.ac.cn) (Based on bash):
```shell
vi ~/.bashrc
export HEPAI_API_KEY2=your_api_key
source ~/.bashrc
```

#### Install and configure from source the DrSai Multi-Agent Collaboration System runtime environment

Create an account on [code.ihep.ac.cn](https://code.ihep.ac.cn/) and clone the OpenDrSai repository locally:
```shell
git clone https://code.ihep.ac.cn/hepai/drsai.git drsai
cd drsai
```

Configure the conda environment and install dependencies:
```shell
conda create -n drsai python=3.10
conda activate drsai
pip install .
```

### 2.2. Deploy the DrSai Multi-Agent Collaboration System backend directly

```python
from DrSai.backend.app_worker import Run_DrSaiAPP, DrSaiModelConfig, DrSaiWorkerConfig
Run_DrSaiAPP(model_args=DrSaiModelConfig, worker_args=DrSaiWorkerConfig).run_drsai()
```     

Save the above code as `app.py`. Running it will start the service on local port 42801:
```shell
python app.py
```

### 2.3. Accessing customized agents using the HepAI client

```python
from hepai import HepAI 
import os
import json
import requests
import sys

HEPAI_API_KEY = os.getenv("HEPAI_API_KEY2")
base_url = "http://localhost:42801/v1"

# Invoke HepAI client interface
client = HepAI(api_key=HEPAI_API_KEY, base_url=base_url)
completion = client.chat.completions.create(
  model='hepai/Dr-Sai',
  messages=[
    {"role": "user", "content": "请使用百度搜索什么是Ptychography?"}
  ],
  stream=True,
  extra_body = {
    "base_models": "openai/gpt-4o-mini",
    "base_url": 'https://aiapi.ihep.ac.cn/apiv2'}
)
for chunk in completion:
  if chunk.choices[0].delta.content:
    print(chunk.choices[0].delta.content, end='', flush=True)
print('\n')
```

## 3. Configure memory layer, executor, and other components

### 3.1 Install and configure [hairag](https://code.ihep.ac.cn/shangzj18/hai-rag-os) as the memory layer component for agents

- Clone hairag and install dependency packages:
```shell
cd .. 
git clone https://code.ihep.ac.cn/shangzj18/hai-rag-os.git hairag
cd hairag
conda create -n hairag python=3.10
conda activate hairag
pip install -r requirements.txt
```

- Download the embedding model to hairag/embedding, refer to [hairag embedding documentation](https://code.ihep.ac.cn/shangzj18/hai-rag-os/-/blob/main/embedding/README.md?ref_type=heads) for details.

- hairag is based on the Qdrant Vector Database framework, see [hairag configuration documentation](https://code.ihep.ac.cn/shangzj18/hai-rag-os/-/tree/main?ref_type=heads) for specific configurations.

For a detailed example, refer to [hairag documentation](https://code.ihep.ac.cn/shangzj18/hai-rag-os/-/blob/main/src/example_rag.ipynb?ref_type=heads).

### 3.2 Install and configure [Dr.Sai Code Worker V2](https://code.ihep.ac.cn/xuliang/drsai-code-worker-v2) as the executor component for agents

- Clone Dr.Sai Code Worker V2 and install dependency packages:
```shell
cd .. 
git clone https://code.ihep.ac.cn/xuliang/drsai-code-worker-v2.git  codeworker
cd codeworker
```
 - Install dependency packages:
```shell
conda create -n codeworker python=3.10
conda activate codeworker
pip install -r requirements.txt
```
For detailed configuration of Dr.Sai Code Worker V2, see [README.md](https://code.ihep.ac.cn/xuliang/drsai-code-worker-v2).

### 3.3 Add memory layer and executor plugins for agents

Inherit DrSai_APP, modify the `default_agents` function in [DrSai/dr_sai.py](DrSai/dr_sai.py) to add memory layer and executor plugins for the specific agent configuration. The default agent list includes `Assistant, Coder, Navigator`:
```python
def default_agents(self, llm_config: Dict, **kwargs) -> List[Type[LearnableAgent]]:
    '''
    Default Agent list:
        - Assistant
        - Coder 
        - Navigator: Queries information from Arxiv/docDB/web, etc.
    When used as a backend, please modify and register your custom Agents here.
    '''
    stream = kwargs.get('stream', True)
    # TODO: Add default RAG
    if self._default_agents is None:
        assistant = AssistantAgent("Assistant", llm_config=llm_config, stream=stream)
        self.register_agent(assistant.name, assistant)
        coder = Coder("Coder", llm_config=llm_config, stream=stream)
        self.register_agent(coder.name, coder)
        navigator = Navigator("Navigator", llm_config=llm_config, stream=stream)
        tool_calls_register.all_tools(agent=navigator)
        self.register_agent(navigator.name, navigator)
        
        self._default_agents = [coder, assistant, navigator]
    return self._default_agents

```
To add a specific synchrotron radiation knowledge base memory layer plugin and an executable code agent for the language model assistant, you can refer to the following configuration:
```python

from DrSai.backend.app_worker import DrSaiAPP
class YourDrSaiAPP(DrSaiAPP):
    
    def __init__(
            self, 
            **kwargs
    ):
        super().__init__(
            **kwargs
        )
    def default_agents(self, llm_config: Dict, **kwargs) -> List[Type[LearnableAgent]]:
        '''
        Custom Agent list:
        '''
        stream = kwargs.get('stream', True)
        if self._default_agents is None:
            # Configure an assistant with synchrotron radiation knowledge
            ## Configure a custom RAG function and corresponding retrieval parameters
            headers={"Authorization": f"Bearer {self.api_key}"}
            request_data = {
                    "model": "hai-rag",
                    "username": "your_username",
                    "method": "retrieve",
                    "similarity_top_k": 10,
                    'collection':"synchrotron",
                    "score_limit": 0.5
                }
            RAG_config = {"request_data": request_data, "headers": headers}
            retrieve_config = {"RAG_function": hai_rag_retrieve, "RAG_config": RAG_config}
            ## Create agent
            assistant = AssistantAgent(
                name="AssistantAgent",
                llm_config=llm_config, # Defines model and API configuration
                retrieve_config = retrieve_config, # Defines RAG configuration and custom RAG function
                system_message = "You are a synchrotron radiation knowledge question-and-answer assistant",
                description="A synchrotron radiation knowledge question-and-answer assistant", 
                stream=stream)
            self.register_agent(assistant.name, assistant)
            # Configure an executable code agent
            tester = AssistantAgent(
                name="AssistantAgent",
                llm_config=llm_config,
                description="A code executor", 
                worker_name={"name": "hepai/code-worker-v2", "base_url": "http://localhost:42899/apiv2"},# Defines worker access configuration
                worker_config={"llm_config": llm_config, "job_timeout": 10}, # Defines worker parameter configuration
                stream=stream)
            self.register_agent(tester.name, tester)
            # Configure the default agent list
            coder = Coder("Coder", llm_config=llm_config, stream=stream)
            self.register_agent(coder.name, coder)
            navigator = Navigator("Navigator", llm_config=llm_config, stream=stream)
            tool_calls_register.all_tools(agent=navigator)
            self.register_agent(navigator.name, navigator)
            
            self._default_agents = [coder, assistant, navigator, tester]
        return self._default_agents

```

### 3.4 Start the DrSai backend service

Save the custom agent YourDrSaiAPP backend service and the following code as run_drsai.py:
```python

from DrSai.backend.app_worker import Run_DrSaiAPP, DrSaiModelConfig, DrSaiWorkerConfig
Run_DrSaiAPP(model_args=DrSaiModelConfig, worker_args=DrSaiWorkerConfig, drsaiapp=YourDrSaiAPP).run_drsai()

```

Start the customized agent backend service:
```shell
python run_drsai.py
```

## 4. Detailed Documentation
See the docs directory:
```shell
.
├── 1-Agent_components.md         # Introduction to agent component development
├── 2-Agent.md                    # Introduction to agent development
├── 3-Multi-Agent_system.md       # Introduction to multi-agent collaborative system development
├── 4-backend.md                  # Introduction to backend service development
├── developer.md                  # Developer's guide
└── 0-Re0:Dr.Sai.md               # Dr.Sai development logic
```

## 5. Contact
For any questions, please contact Dr.Sai team:
- Email：hepai@ihep.ac.cn / xiongdb@ihep.ac.cn
- WeChat：xiongdongbo_12138