Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
384 changes: 384 additions & 0 deletions tests/integration/run_benchmark.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,384 @@
import asyncio
import json
import logging
import os
import shutil
import sys
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from typing import Any, Optional, Type

import pandas as pd
import requests
from pandas import DataFrame
from pydantic import BaseModel

from deepeval.metrics import GEval
from deepeval.models.base_model import DeepEvalBaseLLM
from deepeval.test_case import LLMTestCase, LLMTestCaseParams

from osa_tool.config.settings import ConfigManager
from osa_tool.core.git.git_agent import GitHubAgent, GitLabAgent, GitverseAgent
from osa_tool.core.git.metadata import RepositoryMetadata
from osa_tool.operations.docs.readme_generation.readme_agent import ReadmeAgent
from osa_tool.tools.repository_analysis.sourcerank import SourceRank
from osa_tool.utils.arguments_parser import build_parser_from_yaml
from osa_tool.utils.logger import logger
from osa_tool.utils.utils import delete_repository, format_time, parse_git_url, rich_section

README_QUALITY_CRITERIA = """
Determine whether the AI-generated Readme file (ACTUAL_OUTPUT)
is better than the original one (EXPECTED_OUTPUT).
ACTUAL_OUTPUT contains two fields: 'readme', which contains the generated README itself,
and 'repo_structure' which is json with repository's structure.
Generated README's content must be consistent with the provided repository structure.
The ACTUAL_OUTPUT does not necessary have to be the same as EXPECTED_OUTPUT,
Your goal is to determine which text is better, using the provided Evaluations steps.
Readme structure does not matter much as long as it passes the evaluation steps.
"""

README_QUALITY_STEPS = [
"Step 1: Does the provided structure of the repository address README content?",
"Step 2: Does the README provide a clear and accurate overview of the repository's purpose?",
"Step 3: Are installation and setup instructions included and easy to follow?",
"Step 4: Are usage examples provided and do they clearly demonstrate functionality?",
"Step 5: Are dependencies or requirements listed appropriately?",
"Step 6: Is the README easy to read, well-structured, and free of confusing language?",
]


def _strip_markdown_json_fence(text: str) -> str:
cleaned = text.strip()
if not cleaned.startswith("```"):
return cleaned
cleaned = cleaned.removeprefix("```").strip()
if cleaned.lower().startswith("json"):
cleaned = cleaned[4:].lstrip()
if "```" in cleaned:
cleaned = cleaned.split("```", 1)[0].strip()
return cleaned


class CustomLLM(DeepEvalBaseLLM):
def __init__(
self,
api: str = "openrouter",
model: str = "gpt-4.1",
url: str = "https://openrouter.ai/api/v1",
*,
max_tokens: int = 1024,
request_timeout: float = 180.0,
use_json_object_mode: bool = True,
):
self.api = api
self.model_name = model
self.url = url.rstrip("/")
self.max_tokens = max_tokens
self.request_timeout = request_timeout
self.use_json_object_mode = use_json_object_mode

def load_model(self):
return self

def supports_json_mode(self) -> bool:
return True

def get_model_name(self) -> str:
return self.model_name

def _api_key(self) -> str:
api = (self.api or "").lower().strip()
url = (self.url or "").lower()
openrouter_key = os.getenv("OPENROUTER_API_KEY", "")
openai_key = os.getenv("OPENAI_API_KEY", "")
service_key = os.getenv("LLM_SERVICE_KEY", "")
if api == "openrouter" or "openrouter.ai" in url:
return openrouter_key or openai_key or service_key
if api == "openai":
return openai_key or openrouter_key or service_key
return openrouter_key or openai_key or service_key

def _headers(self) -> dict[str, str]:
key = self._api_key()
if not key:
return {}
headers = {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
}
if "openrouter.ai" in self.url.lower():
headers["HTTP-Referer"] = "https://github.com/aimclub/OSA"
headers["X-Title"] = "OSA README benchmark"
return headers

def _post_chat(self, messages: list[dict[str, str]], *, response_format: Optional[dict[str, str]] = None) -> str:
headers = self._headers()
if not headers:
raise RuntimeError("Missing judge API key. Set OPENROUTER_API_KEY, OPENAI_API_KEY, or LLM_SERVICE_KEY.")
payload: dict[str, Any] = {
"model": self.model_name,
"messages": messages,
"max_tokens": self.max_tokens,
"temperature": 0.0,
}
if response_format and self.use_json_object_mode:
payload["response_format"] = response_format

response = requests.post(
f"{self.url}/chat/completions", headers=headers, json=payload, timeout=self.request_timeout
)
if response.status_code == 200:
return (response.json()["choices"][0]["message"]["content"] or "").strip()

if response_format and self.use_json_object_mode:
payload.pop("response_format", None)
response = requests.post(
f"{self.url}/chat/completions", headers=headers, json=payload, timeout=self.request_timeout
)
if response.status_code == 200:
return (response.json()["choices"][0]["message"]["content"] or "").strip()

raise RuntimeError(f"Judge LLM HTTP {response.status_code}: {response.text[:500]}")

def generate(self, prompt: str) -> str:
return self._post_chat([{"role": "user", "content": prompt}])

async def a_generate(self, prompt: str, schema=None):
return await asyncio.to_thread(self.generate, prompt)


def generate_readme(config_manager: ConfigManager, metadata: RepositoryMetadata, args, safe_name: str) -> str:
readmes_dir = os.path.join(os.path.dirname(args.table_path), "readmes")
os.makedirs(readmes_dir, exist_ok=True)

readme_agent = ReadmeAgent(config_manager=config_manager, metadata=metadata)
dest_path = os.path.join(readmes_dir, f"{safe_name}_README.md")
readme_agent.file_to_save = dest_path

readme_agent.generate_readme()

src = os.path.join(readme_agent.repo_path, "README.md")
if os.path.isfile(src):
shutil.copy2(src, dest_path)
else:
logger.warning(f"README not found at clone path after generation: {src}")

return dest_path


def get_repo_structure_json(repo_path: str) -> str:
"""Gathers a simple repository structure for GEVAL."""
tree = []
for root, dirs, files in os.walk(repo_path):
if ".git" in dirs:
dirs.remove(".git")
rel_path = os.path.relpath(root, repo_path)
tree.append({"dir": rel_path if rel_path != "." else "/", "files": files})
return json.dumps(tree, ensure_ascii=False)


def process_repository(repo_url: str, args) -> dict:
stage_start = time.time()

repos_dir = os.path.join(os.path.dirname(args.table_path), "repositories")
logs_dir = os.path.join(os.path.dirname(args.table_path), "logs")
os.makedirs(logs_dir, exist_ok=True)
os.makedirs(repos_dir, exist_ok=True)

_, _, repo_name, _ = parse_git_url(repo_url)

url_parts = repo_url.rstrip("/").split("/")
safe_name = f"{url_parts[-2]}_{url_parts[-1]}" if len(url_parts) >= 2 else repo_name

worker_dir = os.path.join(repos_dir, safe_name)
os.makedirs(worker_dir, exist_ok=True)

original_cwd = os.getcwd()
os.chdir(worker_dir)

log_file = os.path.join(logs_dir, f"{safe_name}.log")

logger.setLevel(logging.DEBUG)
file_handler = logging.FileHandler(log_file, encoding="utf-8")
file_handler.setLevel(logging.DEBUG)
file_handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(message)s"))
logger.addHandler(file_handler)

result = {"repository": repo_url, "name": repo_name, "status": "Failed", "geval_score": None}

try:
args.repository = repo_url
config_manager = ConfigManager(args)

if not hasattr(config_manager.config, "git"):
config_manager.config.git = type("obj", (object,), {"repository": repo_url})
else:
config_manager.config.git.repository = repo_url

if "github.com" in repo_url:
git_agent = GitHubAgent(repo_url)
elif "gitlab" in repo_url:
git_agent = GitLabAgent(repo_url)
elif "gitverse.ru" in repo_url:
git_agent = GitverseAgent(repo_url)
else:
logger.error(f"Unsupported GIT platform: {repo_url}")
return result

git_agent.clone_repository()

actual_clone_path = os.path.join(worker_dir, repo_name)

expected_output = ""
original_readme_path = os.path.join(actual_clone_path, "README.md")

if not os.path.exists(original_readme_path):
original_readme_path = os.path.join(actual_clone_path, "readme.md")

if os.path.exists(original_readme_path):
with open(original_readme_path, "r", encoding="utf-8", errors="replace") as f:
expected_output = f.read()

repo_structure = get_repo_structure_json(actual_clone_path)

SourceRank(config_manager)
dest_path = generate_readme(config_manager, git_agent.metadata, args, safe_name)

if os.path.exists(dest_path):
result.update({"name": git_agent.metadata.name, "status": "Success"})
logger.info(f"Successfully generated README in {format_time(time.time() - stage_start)}")

logger.info("Starting GEVAL assessment...")
with open(dest_path, "r", encoding="utf-8", errors="replace") as f:
generated_readme = f.read()

judge_model = CustomLLM(api=args.api, model=args.model, url=args.base_url)
metric = GEval(
name="Readme quality",
criteria=README_QUALITY_CRITERIA,
evaluation_steps=README_QUALITY_STEPS,
evaluation_params=[LLMTestCaseParams.ACTUAL_OUTPUT, LLMTestCaseParams.EXPECTED_OUTPUT],
model=judge_model,
verbose_mode=False,
async_mode=False,
)

test_case = LLMTestCase(
input="",
actual_output=json.dumps(
{"readme": generated_readme, "repo_structure": repo_structure}, ensure_ascii=False
),
expected_output=expected_output,
)

try:
metric.measure(test_case)
result["geval_score"] = metric.score
logger.info(f"GEVAL Score: {metric.score}")
except Exception as e:
logger.error(f"GEval metric measurement failed: {e}")

else:
result.update({"name": git_agent.metadata.name, "status": "Failed"})
logger.error(f"Failed to generate README for {git_agent.metadata.name}")

except Exception as e:
logger.error(f"Error processing {repo_url}: {e}")

finally:
file_handler.flush()
file_handler.close()
logger.removeHandler(file_handler)
delete_repository(repo_url)
os.chdir(original_cwd)
shutil.rmtree(worker_dir, ignore_errors=True)

return result


def load_table(table_path: str) -> DataFrame:
if not os.path.isfile(table_path):
test_repos = [
"https://github.com/google/python-fire",
"https://github.com/encode/httpx",
"https://github.com/AntonOsika/gpt-engineer",
"https://github.com/THUDM/ChatGLM-6B",
]

rows = [{"repository": repo, "status": "Pending", "geval_score": None} for repo in test_repos]
df = pd.DataFrame(rows)
df.to_csv(table_path, index=False)
logger.info(f"Created new benchmark run at {table_path} with {len(test_repos)} repos.")
return df

df = pd.read_csv(table_path) if table_path.endswith(".csv") else pd.read_excel(table_path)

if "repository" not in df.columns:
if "repo_url" in df.columns:
df["repository"] = df["repo_url"]
else:
logger.error("Table must contain a 'repository' or 'repo_url' column.")
sys.exit(1)

if "status" not in df.columns:
df["status"] = "Pending"
if "geval_score" not in df.columns:
df["geval_score"] = None

return df


def main():
parser = build_parser_from_yaml(extra_sections=["settings", "arguments", "multi-run"])
args, _ = parser.parse_known_args()

if getattr(args, "table_path", None) is None:
results_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "benchmark_results"))
os.makedirs(results_dir, exist_ok=True)
timestamp = time.strftime("%Y%m%d_%H%M%S")
args.table_path = os.path.join(results_dir, f"run_{timestamp}.csv")

if getattr(args, "api", None) is None:
args.api = "openai"
if getattr(args, "base_url", None) is None:
args.base_url = "https://openrouter.ai/api/v1"
if getattr(args, "model", None) is None:
args.model = "openai/gpt-4.1"

args.table_path = os.path.abspath(args.table_path)

df = load_table(args.table_path)
repos = df["repository"].dropna().tolist()

unprocessed = [r for r in repos if df.loc[df["repository"] == r, "status"].values[0] != "Success"]

if unprocessed:
rich_section(f"Starting lightweight README Generation & GEVAL for {len(unprocessed)} repos")
with ProcessPoolExecutor(max_workers=max(1, os.cpu_count() // 2)) as executor:
futures = {executor.submit(process_repository, repo, args): repo for repo in unprocessed}
for future in as_completed(futures):
repo = futures[future]
try:
res = future.result()
df.loc[df["repository"] == repo, "status"] = res["status"]
df.loc[df["repository"] == repo, "geval_score"] = res.get("geval_score")

if args.table_path.endswith(".csv"):
df.to_csv(args.table_path, index=False)
else:
df.to_excel(args.table_path, index=False)
except Exception as e:
logger.error(f"Failed to process {repo} — {e}")

print("\n" + "=" * 90)
print(" FINAL BENCHMARK RESULTS ".center(90, "="))
print("=" * 90)
print(df.to_string(index=False))
print("=" * 90)
print(f"All files (logs, readmes, table) saved to: {os.path.dirname(args.table_path)}\n")
else:
rich_section("All repositories processed successfully.")


if __name__ == "__main__":
main()
Loading