diff --git a/nodes.py b/nodes.py index f6e769d..e6ed955 100644 --- a/nodes.py +++ b/nodes.py @@ -247,29 +247,32 @@ class AnalyzeRelationships(Node): language = shared.get("language", "english") # Get language use_cache = shared.get("use_cache", True) # Get use_cache flag, default to True + # Get the actual number of abstractions directly + num_abstractions = len(abstractions) + # Create context with abstraction names, indices, descriptions, and relevant file snippets - context = "Identified Abstractions:\n" + context = "Identified Abstractions:\\n" all_relevant_indices = set() abstraction_info_for_prompt = [] for i, abstr in enumerate(abstractions): # Use 'files' which contains indices directly file_indices_str = ", ".join(map(str, abstr["files"])) # Abstraction name and description might be translated already - info_line = f"- Index {i}: {abstr['name']} (Relevant file indices: [{file_indices_str}])\n Description: {abstr['description']}" - context += info_line + "\n" + info_line = f"- Index {i}: {abstr['name']} (Relevant file indices: [{file_indices_str}])\\n Description: {abstr['description']}" + context += info_line + "\\n" abstraction_info_for_prompt.append( f"{i} # {abstr['name']}" ) # Use potentially translated name here too all_relevant_indices.update(abstr["files"]) - context += "\nRelevant File Snippets (Referenced by Index and Path):\n" + context += "\\nRelevant File Snippets (Referenced by Index and Path):\\n" # Get content for relevant files using helper relevant_files_content_map = get_content_for_indices( files_data, sorted(list(all_relevant_indices)) ) # Format file content for context - file_context_str = "\n\n".join( - f"--- File: {idx_path} ---\n{content}" + file_context_str = "\\n\\n".join( + f"--- File: {idx_path} ---\\n{content}" for idx_path, content in relevant_files_content_map.items() ) context += file_context_str @@ -277,15 +280,21 @@ class AnalyzeRelationships(Node): return ( context, "\n".join(abstraction_info_for_prompt), + num_abstractions, # Pass the actual count project_name, language, use_cache, ) # Return use_cache def exec(self, prep_res): - context, abstraction_listing, project_name, language, use_cache = ( - prep_res # Unpack use_cache - ) + ( + context, + abstraction_listing, + num_abstractions, # Receive the actual count + project_name, + language, + use_cache, + ) = prep_res # Unpack use_cache print(f"Analyzing relationships using LLM...") # Add language instruction and hints only if not English @@ -335,7 +344,7 @@ relationships: Now, provide the YAML output: """ - response = call_llm(prompt) + response = call_llm(prompt, use_cache=use_cache) # --- Validation --- yaml_str = response.strip().split("```yaml")[1].split("```")[0].strip() @@ -354,7 +363,6 @@ Now, provide the YAML output: # Validate relationships structure validated_relationships = [] - num_abstractions = len(abstraction_listing.split("\n")) for rel in relationships_data["relationships"]: # Check for 'label' key if not isinstance(rel, dict) or not all( diff --git a/utils/call_llm.py b/utils/call_llm.py index 714d4dd..6e13975 100644 --- a/utils/call_llm.py +++ b/utils/call_llm.py @@ -47,7 +47,7 @@ def call_llm(prompt: str, use_cache: bool = True) -> str: logger.info(f"RESPONSE: {cache[prompt]}") return cache[prompt] - # Call the LLM if not in cache or cache disabled + # # Call the LLM if not in cache or cache disabled # client = genai.Client( # vertexai=True, # # TODO: change to your own project id and location @@ -59,8 +59,8 @@ def call_llm(prompt: str, use_cache: bool = True) -> str: client = genai.Client( api_key=os.getenv("GEMINI_API_KEY", ""), ) - # model = os.getenv("GEMINI_MODEL", "gemini-2.5-pro-exp-03-25") - model = os.getenv("GEMINI_MODEL", "gemini-2.0-flash-exp") + model = os.getenv("GEMINI_MODEL", "gemini-2.5-pro-exp-03-25") + # model = os.getenv("GEMINI_MODEL", "gemini-2.5-flash-preview-04-17") response = client.models.generate_content(model=model, contents=[prompt]) response_text = response.text