redcodegen 0.0.3__tar.gz → 0.0.5__tar.gz

{redcodegen-0.0.3 → redcodegen-0.0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: redcodegen
-Version: 0.0.3
+Version: 0.0.5
 Summary: Add your description here
 Requires-Dist: click>=8.0.0
 Requires-Dist: cwe2>=3.0.0

{redcodegen-0.0.3 → redcodegen-0.0.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "redcodegen"
-version = "0.0.3"
+version = "0.0.5"
 description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.11"

{redcodegen-0.0.3 → redcodegen-0.0.5}/redcodegen/constants.py

@@ -11,7 +11,7 @@ CWE_TOP_25 = [
     798, 190, 400, 306
 ]
 
-def create_lm(model_name="openai/gpt-4o-mini", temperature=0.8, api_key=None):
+def create_lm(model_name="openai/gpt-4o-mini", temperature=0.8, api_key=None, api_base=None):
     """Create a DSPy language model instance.
 
     Args:
@@ -23,12 +23,24 @@ def create_lm(model_name="openai/gpt-4o-mini", temperature=0.8, api_key=None):
     """
     if api_key is None:
         api_key = os.environ.get("OPENAI_API_KEY")
-    return dspy.LM(
-        model_name,
-        api_key=api_key,
-        temperature=temperature,
-        max_tokens=16000
-    )
+    if api_base is None:
+        api_base = os.environ.get("OPENAI_API_BASE")
+
+    if api_base is None:
+        return dspy.LM(
+            model_name,
+            api_key=api_key,
+            temperature=temperature,
+            max_tokens=16000
+        )
+    else:
+        return dspy.LM(
+            model_name,
+            api_key=api_key,
+            api_base=api_base,
+            temperature=temperature,
+            max_tokens=16000
+        )
 
 SCENARIO_EXAMPLES = [
     dspy.Example(

{redcodegen-0.0.3 → redcodegen-0.0.5}/redcodegen/main.py

@@ -155,13 +155,18 @@ def append_to_jsonl(record: Dict[str, Any], output_path: Path):
     default=None,
     help='API key (defaults to OPENAI_API_KEY env var)'
 )
+@click.option(
+    '--api-base',
+    default=None,
+    help='API base URL (defaults to OPENAI_API_BASE env var)'
+)
 @click.option(
     '--temperature',
     default=0.8,
     type=float,
     help='Temperature for code generation (default: 0.8)'
 )
-def main(cwes, use_top_25, min_samples, output, model, api_key, temperature):
+def main(cwes, use_top_25, min_samples, output, model, api_key, api_base, temperature):
     """Generate and evaluate vulnerable code samples for specified CWEs.
 
     Examples:
@@ -172,7 +177,7 @@ def main(cwes, use_top_25, min_samples, output, model, api_key, temperature):
         python -m redcodegen --use-top-25 --model openai/gpt-4o # switch model
     """
     # Configure DSPy with specified model
-    lm = create_lm(model_name=model, temperature=temperature, api_key=api_key)
+    lm = create_lm(model_name=model, temperature=temperature, api_key=api_key, api_base=api_base)
     dspy.configure(lm=lm)
     logger.info(f"Configured model: {model}")
 

{redcodegen-0.0.3 → redcodegen-0.0.5}/redcodegen/scenarios.py

@@ -29,8 +29,8 @@ class SuggestLibraries(dspy.Signature):
     task: str = dspy.InputField()
     suggested_libraries: List[str] = dspy.InputField()
 
-    chosen_library: Optional[str] = dspy.OutputField(desc="choose a library that would best help solve the task, or None")
-    rephrased_task: Optional[str] = dspy.OutputField(desc="rephrase the task in terms of the chosen library, or None")
+    chosen_library: str = dspy.OutputField(desc="choose a library that would best help solve the task, or say None")
+    rephrased_task: str = dspy.OutputField(desc="rephrase the task in terms of the chosen library, or say None")
 suggest_libraries = dspy.Predict(SuggestLibraries)
 
 def generate(cwe_id, min_scenarios=3):
@@ -48,12 +48,12 @@ def generate(cwe_id, min_scenarios=3):
     output_scenarios = []
     while len(output_scenarios) < min_scenarios:
         scenarios = extract_scenarios(name=entry.name, description=entry.extended_description,
-                                      config={"temperature": 0.8, "rollout_id": len(output_scenarios)}).scenarios
+                                      config={"rollout_id": len(output_scenarios)}).scenarios
         output_scenarios.extend(scenarios)
     scenarios = [strip_vulnerability(scenario=i).coding_task for i in output_scenarios]
     suggestions = [suggest_libraries(task=i, suggested_libraries=CODEQL_LIBRARIES) for i in scenarios]
     results = [
-        i.rephrased_task if i.rephrased_task is not None else j
+        i.rephrased_task if ((i.rephrased_task is not None) and (i.rephrased_task.lower().strip() != "none")) else j
         for i,j in zip(suggestions, scenarios)
     ]