Merge pull request #7 from AMD-AGI/add_pip_package

Convert GEAK-Agent to installable pip package

Author: jianghui-jianghui

.gitignore

@@ -1,3 +1,21 @@
-src/__pycache__/
-src/*/__pycache__/
-src/*/*/__pycache__/
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.egg-info/
+dist/
+build/
+*.egg
+
+# Virtual environments
+venv/
+.venv/
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+
+# Jupyter
+.ipynb_checkpoints/

geak_agent/__init__.py

@@ -0,0 +1,82 @@
+# Copyright(C) [2025] Advanced Micro Devices, Inc. All rights reserved.
+
+"""
+GEAK-Agent: An LLM-based multi-agent framework for generating functional and efficient GPU kernels.
+"""
+
+__version__ = "0.1.0"
+
+# Lazy imports - only import when accessed to avoid loading tb_eval dependency at import time
+__all__ = [
+    # Version
+    "__version__",
+    # Agents
+    "BaseAgent",
+    "SequentialBaseAgent",
+    "GaAgent",
+    "Reflexion",
+    "Reflexion_Oneshot",
+    "OptimAgent",
+    "DirectPrompt",
+    # Models
+    "BaseModel",
+    "OpenAIModel",
+    "StandardOpenAIModel",
+    "ClaudeModel",
+    "StandardClaudeModel",
+    "GeminiModel",
+    # Dataloaders
+    "TritonBench",
+    "ROCm",
+    "ProblemState",
+    "ProblemStateROCm",
+    # Utils
+    "load_config",
+]
+
+
+def __getattr__(name):
+    """Lazy import to avoid loading tb_eval at package import time."""
+    if name in ("BaseAgent", "SequentialBaseAgent"):
+        from geak_agent.agents.Base import BaseAgent, SequentialBaseAgent
+        return BaseAgent if name == "BaseAgent" else SequentialBaseAgent
+    elif name == "GaAgent":
+        from geak_agent.agents.GaAgent import GaAgent
+        return GaAgent
+    elif name == "Reflexion":
+        from geak_agent.agents.Reflexion import Reflexion
+        return Reflexion
+    elif name == "Reflexion_Oneshot":
+        from geak_agent.agents.reflexion_oneshot import Reflexion_Oneshot
+        return Reflexion_Oneshot
+    elif name == "OptimAgent":
+        from geak_agent.agents.OptimAgent import OptimAgent
+        return OptimAgent
+    elif name == "DirectPrompt":
+        from geak_agent.agents.DirectPrompt import DirectPrompt
+        return DirectPrompt
+    elif name == "BaseModel":
+        from geak_agent.models.Base import BaseModel
+        return BaseModel
+    elif name in ("OpenAIModel", "StandardOpenAIModel"):
+        from geak_agent.models.OpenAI import OpenAIModel, StandardOpenAIModel
+        return OpenAIModel if name == "OpenAIModel" else StandardOpenAIModel
+    elif name in ("ClaudeModel", "StandardClaudeModel"):
+        from geak_agent.models.Claude import ClaudeModel, StandardClaudeModel
+        return ClaudeModel if name == "ClaudeModel" else StandardClaudeModel
+    elif name == "GeminiModel":
+        from geak_agent.models.Gemini import GeminiModel
+        return GeminiModel
+    elif name == "TritonBench":
+        from geak_agent.dataloaders.TritonBench import TritonBench
+        return TritonBench
+    elif name == "ROCm":
+        from geak_agent.dataloaders.ROCm import ROCm
+        return ROCm
+    elif name in ("ProblemState", "ProblemStateROCm"):
+        from geak_agent.dataloaders.ProblemState import ProblemState, ProblemStateROCm
+        return ProblemState if name == "ProblemState" else ProblemStateROCm
+    elif name == "load_config":
+        from geak_agent.args_config import load_config
+        return load_config
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

src/agents/Base.py geak_agent/agents/Base.pysrc/agents/Base.py renamed to geak_agent/agents/Base.py

@@ -6,9 +6,9 @@
 from loguru import logger
 from dataclasses import dataclass, asdict
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from models.Base import BaseModel
-from dataloaders.ProblemState import ProblemState
-from memories.Memory import BaseMemory
+from geak_agent.models.Base import BaseModel
+from geak_agent.dataloaders.ProblemState import ProblemState
+from geak_agent.memories.Memory import BaseMemory
 
 
 

src/agents/DirectPrompt.py geak_agent/agents/DirectPrompt.pysrc/agents/DirectPrompt.py renamed to geak_agent/agents/DirectPrompt.py

@@ -1,7 +1,7 @@
 # Copyright(C) [2025] Advanced Micro Devices, Inc. All rights reserved.
 
-from agents.Base import BaseAgent
-from utils.utils import clear_code
+from geak_agent.agents.Base import BaseAgent
+from geak_agent.utils.utils import clear_code
 
 class DirectPrompt(BaseAgent):
     def run_single_pass(self, mem, verbose=False):

src/agents/GaAgent.py geak_agent/agents/GaAgent.pysrc/agents/GaAgent.py renamed to geak_agent/agents/GaAgent.py

@@ -4,13 +4,13 @@
 import os
 import json
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from agents.reflexion_oneshot import Reflexion_Oneshot
-from utils.utils import clear_code, extract_function_signatures, clear_json
-from memories.Memory import MemoryClassMeta
-from prompts import prompt_for_generation, prompt_for_reflection
+from geak_agent.agents.reflexion_oneshot import Reflexion_Oneshot
+from geak_agent.utils.utils import clear_code, extract_function_signatures, clear_json
+from geak_agent.memories.Memory import MemoryClassMeta
+from geak_agent.prompts import prompt_for_generation, prompt_for_reflection
 from loguru import logger
 from tenacity import RetryError
-from dataloaders.ProblemState import tempCode
+from geak_agent.dataloaders.ProblemState import tempCode
 from typing import List, Optional
 
 class GaAgent(Reflexion_Oneshot):

src/agents/GaAgent_ROCm.py geak_agent/agents/GaAgent_ROCm.pysrc/agents/GaAgent_ROCm.py renamed to geak_agent/agents/GaAgent_ROCm.py

@@ -4,13 +4,13 @@
 import os
 import json
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from agents.reflexion_oneshot import Reflexion_Oneshot
-from utils.utils import clear_code, extract_function_signatures, clear_json
-from memories.Memory import MemoryClassMeta
-from prompts import prompt_for_generation, prompt_for_reflection
+from geak_agent.agents.reflexion_oneshot import Reflexion_Oneshot
+from geak_agent.utils.utils import clear_code, extract_function_signatures, clear_json
+from geak_agent.memories.Memory import MemoryClassMeta
+from geak_agent.prompts import prompt_for_generation, prompt_for_reflection
 from loguru import logger
 from tenacity import RetryError
-from dataloaders.ProblemState import tempCode
+from geak_agent.dataloaders.ProblemState import tempCode
 from typing import List, Optional
 from tb_eval.evaluators.interface import get_evaluators
 class GaAgent(Reflexion_Oneshot):

src/agents/OptimAgent.py geak_agent/agents/OptimAgent.pysrc/agents/OptimAgent.py renamed to geak_agent/agents/OptimAgent.py

@@ -4,10 +4,10 @@
 import os
 import json
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from agents.reflexion_oneshot import Reflexion_Oneshot
-from utils.utils import clear_code, extract_function_signatures, clear_json
-from memories.Memory import MemoryClassMeta
-from prompts import prompt_for_generation, prompt_for_reflection
+from geak_agent.agents.reflexion_oneshot import Reflexion_Oneshot
+from geak_agent.utils.utils import clear_code, extract_function_signatures, clear_json
+from geak_agent.memories.Memory import MemoryClassMeta
+from geak_agent.prompts import prompt_for_generation, prompt_for_reflection
 from loguru import logger
 from tenacity import RetryError
 

src/agents/Reflexion.py geak_agent/agents/Reflexion.pysrc/agents/Reflexion.py renamed to geak_agent/agents/Reflexion.py

@@ -1,10 +1,10 @@
 # Copyright(C) [2025] Advanced Micro Devices, Inc. All rights reserved.
 
-from agents.Base import SequentialBaseAgent, BaseAgent
-from utils.utils import clear_code
-from prompts import prompt_for_reflection
-from memories.Memory import ReflexionMemory
-from models.Base import BaseModel
+from geak_agent.agents.Base import SequentialBaseAgent, BaseAgent
+from geak_agent.utils.utils import clear_code
+from geak_agent.prompts import prompt_for_reflection
+from geak_agent.memories.Memory import ReflexionMemory
+from geak_agent.models.Base import BaseModel
 
 
 

geak_agent/agents/__init__.py

@@ -0,0 +1,21 @@
+# Copyright(C) [2025] Advanced Micro Devices, Inc. All rights reserved.
+
+from geak_agent.agents.Base import BaseAgent, SequentialBaseAgent
+from geak_agent.agents.GaAgent import GaAgent
+from geak_agent.agents.GaAgent_ROCm import GaAgent as GaAgentROCm
+from geak_agent.agents.Reflexion import Reflexion
+from geak_agent.agents.reflexion_oneshot import Reflexion_Oneshot
+from geak_agent.agents.OptimAgent import OptimAgent
+from geak_agent.agents.DirectPrompt import DirectPrompt
+
+__all__ = [
+    "BaseAgent",
+    "SequentialBaseAgent",
+    "GaAgent",
+    "GaAgentROCm",
+    "Reflexion",
+    "Reflexion_Oneshot",
+    "OptimAgent",
+    "DirectPrompt",
+]
+

src/agents/reflexion_oneshot.py geak_agent/agents/reflexion_oneshot.pysrc/agents/reflexion_oneshot.py renamed to geak_agent/agents/reflexion_oneshot.py

@@ -5,13 +5,13 @@
 from loguru import logger
 import json
 from dataclasses import asdict
-from agents.Reflexion import Reflexion
-from utils.utils import extract_function_signatures, clear_code, extract_function_calls
-from prompts import prompt_for_reflection
-from memories.Memory import MemoryClassMeta
-from models.Base import BaseModel
-from retrievers.retriever import BM25Retriever
-from prompts import prompt_for_generation
+from geak_agent.agents.Reflexion import Reflexion
+from geak_agent.utils.utils import extract_function_signatures, clear_code, extract_function_calls
+from geak_agent.prompts import prompt_for_reflection
+from geak_agent.memories.Memory import MemoryClassMeta
+from geak_agent.models.Base import BaseModel
+from geak_agent.retrievers.retriever import BM25Retriever
+from geak_agent.prompts import prompt_for_generation
 from concurrent.futures import ThreadPoolExecutor, as_completed