add the creating of google adk template

Author: root

src/abacusagent/create_template.py

@@ -0,0 +1,13 @@
+from pathlib import Path
+from typing import Union, Optional
+import os, sys
+
+def create_google_adk_template(path: Optional[str | Path] = "."):
+    tpath = os.path.join(path, "abacus-agent")
+    os.makedirs(tpath, exist_ok=True)
+    current_file_path = Path(__file__).parent / "google-adk-agent-template.py"
+    
+    # copy the template file to the target directory    
+    os.system(f"cp {current_file_path} {tpath}/agent.py")
+    with open(os.path.join(tpath, "__init__.py"), "w") as file:
+        file.write("from . import agent")

src/abacusagent/env.py

@@ -11,6 +11,11 @@
     "ABACUSAGENT_HOST": "localhost",
     "ABACUSAGENT_PORT": "50001", 
     "ABACUSAGENT_MODEL": "dp",  # fastmcp, abacus, dp
+    
+    # LLM settings
+    "LLM_MODEL": "",
+    "LLM_API_KEY": "",
+    "LLM_BASE_URL": "",
 
     # bohrium settings
     "BOHRIUM_USERNAME": "",
@@ -32,6 +37,9 @@
         "ABACUSAGENT_HOST": "The host address for the AbacusAgent server.",
         "ABACUSAGENT_PORT": "The port number for the AbacusAgent server.",
         "ABACUSAGENT_MODEL": "The model to use for AbacusAgent, can be 'fastmcp', 'test', or 'dp'.",
+        "LLM_MODEL": "The model name for the LLM to use. Like: openai/qwen-turbo, deepseek/deepseek-chat",
+        "LLM_API_KEY": "The API key for the LLM service.",
+        "LLM_BASE_URL": "The base URL for the LLM service, if applicable.",
         "BOHRIUM_USERNAME": "The username for Bohrium.",        
         "BOHRIUM_PASSWORD": "The password for Bohrium.",
         "BOHRIUM_PROJECT_ID": "The project ID for Bohrium.",

src/abacusagent/google-adk-agent-template.py

@@ -0,0 +1,88 @@
+from google.adk.agents import Agent
+from google.adk.models.lite_llm import LiteLlm
+from dp.agent.adapter.adk import CalculationMCPToolset
+from google.adk.tools.mcp_tool.mcp_session_manager import SseServerParams, StreamableHTTPServerParams
+
+import os, json
+
+# Set the secret key in ~/.abacusagent/env.json or as an environment variable, or modify the code to set it directly.
+env_file = os.path.expanduser("~/.abacusagent/env.json")
+if os.path.isfile(env_file):
+    env = json.load(open(env_file, "r"))
+else:
+    env = {}
+model_name = env.get("LLM_MODEL", os.environ.get("LLM_MODEL", ""))
+model_api_key = env.get("LLM_API_KEY", os.environ.get("LLM_API_KEY", ""))
+model_base_url = env.get("LLM_BASE_URL", os.environ.get("LLM_BASE_URL", ""))
+bohrium_username = env.get("BOHRIUM_USERNAME", os.environ.get("BOHRIUM_USERNAME", ""))
+bohrium_password = env.get("BOHRIUM_PASSWORD", os.environ.get("BOHRIUM_PASSWORD", ""))
+bohrium_project_id = env.get("BOHRIUM_PROJECT_ID", os.environ.get("BOHRIUM_PROJECT_ID", ""))
+
+instruction = """You are an expert in materials science and computational chemistry. "
+                "Help users perform ABACUS including single point calculation, structure optimization, molecular dynamics and property calculations. "
+                "The website of ABACUS documentation is at https://abacus.deepmodeling.com/en/latest/, please read it if necessary." 
+                "Use default parameters if the users do not mention, but let users confirm them before submission. "
+                "Always prepare an directory containing ABACUS input files before use specific tool functions."
+                "Always verify the input parameters to users and provide clear explanations of results."
+                "Do not try to modify the input files without explicit permission when errors occured."
+                "The LCAO basis is prefered."
+"""
+
+executor = {
+    "bohr": {
+        "type": "dispatcher",
+        "machine": {
+            "batch_type": "Bohrium",
+            "context_type": "Bohrium",
+            "remote_profile": {
+                "email": bohrium_username,
+                "password": bohrium_password,
+                "program_id": bohrium_project_id,
+                "input_data": {
+                    "image_name": "registry.dp.tech/dptech/dp/native/prod-22618/abacus-agent-tools:v0.0.3-20250703",
+                    "job_type": "container",
+                    "platform": "ali",
+                    "scass_type": "c32_m64_cpu",
+                },
+            },
+        }
+    },
+    "local": {"type": "local",}
+}
+
+EXECUTOR_MAP = {
+    "run_abacus_onejob": executor["bohr"],
+    "abacus_prepare": executor["local"],
+    "generate_bulk_structure": executor["local"],
+    "generate_molecule_structure": executor["local"],
+}
+
+toolset = CalculationMCPToolset(
+    connection_params=SseServerParams(
+        url="http://localhost:50001/sse", # Or any other MCP server URL
+        sse_read_timeout=3000,  # Set SSE timeout to 3000 seconds
+    ),
+    executor_map = EXECUTOR_MAP,
+    executor=executor["local"],
+    storage={
+        "type": "bohrium",
+        "username": bohrium_username,
+        "password": bohrium_password,
+        "project_id": bohrium_project_id,
+    },
+    
+)
+
+root_agent = Agent(
+    name='agent',
+    model=LiteLlm(
+        model=model_name,
+        api_base=model_base_url,
+        api_key=model_api_key
+    ),
+    description=(
+        "Do ABACUS calculations."
+    ),
+    instruction=instruction,
+    tools=[toolset]
+)

src/abacusagent/main.py

@@ -56,6 +56,14 @@ def parse_args():
         default=None,
         help="Host to run the MCP server on (default: localhost)"
     )
+    parser.add_argument(
+        "--create",
+        type=str,
+        nargs='?',
+        default=None,
+        const=".",
+        help="Create a template for Google ADK agent in the specified directory (default: current directory)"
+    )
 
     args = parser.parse_args()
 
@@ -93,6 +101,13 @@ def main():
         model_input=args.model,
         port_input=args.port, 
         host_input=args.host)
+    
+    if args.create is not None:
+        from abacusagent.create_template import create_google_adk_template
+        create_google_adk_template(args.create)
+        print(f"Google ADK agent template created at {args.create}/abacus-agent/agent.py")
+        return
+    
     create_workpath()
 
     from abacusagent.init_mcp import mcp