diff --git a/agent/llm/groq_client.py b/agent/llm/groq_client.py index 9868b8f..b279db4 100644 --- a/agent/llm/groq_client.py +++ b/agent/llm/groq_client.py @@ -48,6 +48,38 @@ def generate_review(diff: str,context : list = None) -> str: except Exception as e: return f"Error generating review: {str(e)}" +def infer_intent(diff: str) -> dict: + from agent.llm.prompts import INTENT_PROMPT + import json + + prompt = INTENT_PROMPT.format(diff=diff[:8000]) + try: + response = client.chat.completions.create( + model = model, + messages = [ + { + "role": "user", + "content": prompt + } + ] + ) + content = response.choices[0].message.content.strip() + + #cleaning the markdown if added by model + if content.startswith("```"): + content = content.split("```")[1] # Extract code from markdown + if content.startswith("json"): + content = content[len("json"):] # Remove language specifier + return json.loads(content) + except Exception as e: + return { + "purpose": "", + "properties": [], + "edge_cases": [], + "error": str(e) + + } + #local test for generate review function if __name__ == "__main__": test_diff = """ @@ -62,4 +94,5 @@ def generate_review(diff: str,context : list = None) -> str: ] result = generate_review(test_diff, context=test_context) #added conetxt parameter (kindly see @pleasingsunlight) - print(result) \ No newline at end of file + print(result) + diff --git a/agent/llm/prompts.py b/agent/llm/prompts.py index 1365eff..c5b08ce 100644 --- a/agent/llm/prompts.py +++ b/agent/llm/prompts.py @@ -17,4 +17,27 @@ - Do NOT repeat the diff - Use context if helpful - If everything looks good, say: "Code looks good ✅" +""" +INTENT_PROMPT = """ +Your are a senior software engineer. + +Analyze the given code diff and extract: + +1. Purpose of the code +2. Key properties (invariants) +3. Edge cases + +STRICT RULES: +- Output only in JSON +- No explanations text +- Keep it concise + +FORMAT: +{"purpose": "...", + "properties": ["...", "..."], + "edge_cases": ["...", "..."] +} + +---DIFF--- +{diff} """ \ No newline at end of file diff --git a/agent/llm/test_generator.py b/agent/llm/test_generator.py index 733bba7..0701db4 100644 --- a/agent/llm/test_generator.py +++ b/agent/llm/test_generator.py @@ -1,45 +1,71 @@ from agent.llm.groq_client import client, model -def generate_tests(diff:str,context:list =None)->str: + +def generate_tests(diff: str, context: list = None, intent: dict = None) -> str: + context_text = "" if context: - context_text = "\n\n".join(context[:5]) # Include only the first 5 chunks for context + context_text = "\n\n".join(context[:5]) + + # Preparing intent text + intent_text = "" + if intent: + properties = ", ".join(intent.get("properties", [])) + edge_cases = ", ".join(intent.get("edge_cases", [])) + + intent_text = f""" +Purpose: +{intent.get("purpose", "")} + +Properties: +{properties} + +Edge Cases: +{edge_cases} +""" + + prompt = f""" - You are a senior software engineer. - Generate concise pytest test cases for the given code diff. +You are a senior software engineer. - STRICT RULES: - - Output only code - - Max 5 test functions - - keep tests short - - Focus on edge cases and core logic - - NO explanations, NO extra text +Generate pytest test cases using: +- Code diff +- Repository context +- Function intent +STRICT RULES: +- Output ONLY Python code +- Max 5 test functions +- Keep tests short +- Focus on edge cases and properties +- NO explanations +- NO markdown (no ```) - ---CONTEXT--- - {context_text} +---INTENT--- +{intent_text} - ---DIFF--- - {diff} - """ +---CONTEXT--- +{context_text} +---DIFF--- +{diff} +""" try: response = client.chat.completions.create( - model = model, - messages = [ - - { - "role": "user", - "content": prompt - } + model=model, + messages=[ + {"role": "user", "content": prompt} ], - ) - content =response.choices[0].message.content.strip() + ) + + content = response.choices[0].message.content.strip() + #cleaning the markdown if added by model if content.startswith("```"): - content = content.split("```")[1] # Extract code from markdown + content = content.split("```")[1] if content.startswith("python"): - content = content[len("python"):] # Remove language specifier + content = content[len("python"):] + return content.strip() + except Exception as e: - return f"Error generating tests:{str(e)}" - \ No newline at end of file + return f"Error generating tests: {str(e)}" \ No newline at end of file diff --git a/tests/test_generated.py b/tests/test_generated.py index ff04abe..a96f787 100644 --- a/tests/test_generated.py +++ b/tests/test_generated.py @@ -1,26 +1,28 @@ -import pytest -import os -from agent.github.committer import commit_tests -from agent.main import get_pr_diff +def test_infer_intent_empty_diff(): + assert infer_intent("") == {"purpose": "", "properties": [], "edge_cases": [], "error": ""} -def test_commit_tests(): - commit_tests() +def test_infer_intent_short_diff(): + diff = "diff --git a/file.py b/file.py" + intent = infer_intent(diff) + assert "purpose" in intent + assert "properties" in intent + assert "edge_cases" in intent -def test_get_pr_diff_empty(): - with pytest.raises(subprocess.CalledProcessError): - get_pr_diff() +def test_infer_intent_long_diff(): + diff = "diff --git a/file.py b/file.py" * 10000 + intent = infer_intent(diff) + assert "purpose" in intent + assert "properties" in intent + assert "edge_cases" in intent -def test_commit_tests_exception(): - try: - commit_tests() - except Exception as e: - assert str(e) +def test_generate_tests_no_context_no_intent(): + diff = "diff --git a/file.py b/file.py" + tests = generate_tests(diff) + assert tests.strip() != "" -def test_get_pr_diff_no_diff(): - diff = get_pr_diff() - assert diff.strip() == "" - -def test_commit_tests_push(): - subprocess.run(["git", "config", "user.name", "github-actions"], check=True) - subprocess.run(["git", "config", "user.email", "actions@github.com"], check=True) - commit_tests() \ No newline at end of file +def test_generate_tests_with_context_and_intent(): + diff = "diff --git a/file.py b/file.py" + context = ["context1", "context2"] + intent = {"purpose": "test", "properties": ["prop1", "prop2"], "edge_cases": ["case1", "case2"]} + tests = generate_tests(diff, context, intent) + assert tests.strip() != "" \ No newline at end of file