anulax
/
doc-gen
1
0
Fork 0
Code Issues Pull Requests Packages project_board Releases Wiki Activity

Starting tool completion

Browse Source
main
Vinayak Ambigapathy ago%!(EXTRA string=1 week)
parent b2aa5dd6a6
commit 585be3db10
1 changed files with 85 additions and 64 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Unified View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 141
      doc_think.py

141
doc_think.py
Unescape View File

@ -149,17 +149,38 @@ class OllamaDocGenerator:
"content": prompt "content": prompt
} }
], ],
"tools": [
{
'type': 'function',
'function': {
'name': 'analyze_file',
'description': 'This tool allows you to examine other Python files in the project and it returns the same structured information you received for the current file (imports, classes, functions, constants, etc.).',
'parameters': {
'type': 'object',
'properties': {
'path': {
'type': 'string',
'description': 'Relative (from the root of the current project) path to the file',
},
},
'required': ['path'],
},
},
},
],
"stream": False, "stream": False,
"options": { "options": {
"temperature": 0.1, "temperature": 0.1,
"top_p": 0.9, "top_p": 0.9,
} }
}, },
timeout=600 # 10 minute timeout for thinking models timeout=60 * 60 * 24
) )
if response.status_code == 200: if response.status_code == 200:
result = response.json() result = response.json()
tool_calls = result.get('tool_calls', [])
print(result, tool_calls)
message = result.get('message', {}) message = result.get('message', {})
content = message.get('content', '') content = message.get('content', '')
# Parse and display thinking process # Parse and display thinking process
@ -171,11 +192,10 @@ class OllamaDocGenerator:
return final_answer if final_answer else content return final_answer if final_answer else content
else: else:
print(f"Error generating documentation: {response.status_code}") print(f"Error generating documentation: {response.status_code} {response.text}")
return None return None
else: else:
print("None thinking model chosen") print("None thinking model chosen")
# Standard generation for regular models
response = self.session.post( response = self.session.post(
f"{self.ollama_url}/api/generate", f"{self.ollama_url}/api/generate",
json={ json={
@ -188,7 +208,7 @@ class OllamaDocGenerator:
"top_p": 0.9, "top_p": 0.9,
} }
}, },
timeout=300 # 5 minute timeout timeout=60 * 60 * 12
) )
if response.status_code == 200: if response.status_code == 200:
@ -205,7 +225,6 @@ class OllamaDocGenerator:
"""Parse thinking model response to extract thinking process and final answer.""" """Parse thinking model response to extract thinking process and final answer."""
import re import re
# Try different thinking tag patterns
thinking_patterns = [ thinking_patterns = [
r'<thinking>(.*?)</thinking>', r'<thinking>(.*?)</thinking>',
r'<think>(.*?)</think>', r'<think>(.*?)</think>',
@ -220,13 +239,10 @@ class OllamaDocGenerator:
match = re.search(pattern, content, re.DOTALL) match = re.search(pattern, content, re.DOTALL)
if match: if match:
thinking_content = match.group(1).strip() thinking_content = match.group(1).strip()
# Remove thinking section from final answer
final_answer = re.sub(pattern, '', content, flags=re.DOTALL).strip() final_answer = re.sub(pattern, '', content, flags=re.DOTALL).strip()
break break
# If no thinking tags found, check for other patterns like "I need to think about..."
if not thinking_content: if not thinking_content:
# Look for thinking indicators at the start
thinking_indicators = [ thinking_indicators = [
r'^(Let me think about.*?(?=\n\n|\n#|\nI\'ll))', r'^(Let me think about.*?(?=\n\n|\n#|\nI\'ll))',
r'^(I need to analyze.*?(?=\n\n|\n#|\nI\'ll))', r'^(I need to analyze.*?(?=\n\n|\n#|\nI\'ll))',
@ -243,15 +259,10 @@ class OllamaDocGenerator:
return thinking_content, final_answer return thinking_content, final_answer
def create_documentation_prompt(self, file_analysis: Dict, project_context: Dict) -> str: def get_code_documentation(self, file_analysis: Dict, project_context: Dict):
"""Create a comprehensive prompt for documentation generation."""
file_path = file_analysis['file_path'] file_path = file_analysis['file_path']
relative_path = os.path.relpath(file_path, project_context['root_path']) relative_path = os.path.relpath(file_path, project_context['root_path'])
return f"""## PROJECT CONTEXT:
prompt = f"""You are a technical documentation expert. Generate comprehensive markdown documentation for the Python file: `{relative_path}`
## PROJECT CONTEXT:
- **Project Root**: {project_context['root_path']} - **Project Root**: {project_context['root_path']}
- **Total Python Files**: {len(project_context['all_files'])} - **Total Python Files**: {len(project_context['all_files'])}
- **External Dependencies**: {', '.join(project_context['external_dependencies']) if project_context['external_dependencies'] else 'None detected'} - **External Dependencies**: {', '.join(project_context['external_dependencies']) if project_context['external_dependencies'] else 'None detected'}
@ -283,10 +294,27 @@ class OllamaDocGenerator:
{file_analysis['content']} {file_analysis['content']}
``` ```
## DOCUMENTATION REQUIREMENTS: """
def create_documentation_prompt(self, file_analysis: Dict, project_context: Dict) -> str:
file_path = file_analysis['file_path']
relative_path = os.path.relpath(file_path, project_context['root_path'])
prompt = f"""You are a technical documentation expert. Generate comprehensive markdown documentation for the Python file: `{relative_path}`
## AVAILABLE TOOLS
Generate a complete markdown documentation file that includes: You have access to a `analyze_file` tool that allows you to examine other Python files in the project. This tool returns the same structured information you received for the current file (imports, classes, functions, constants, etc.).
**When to use this tool:**
- When you need to understand how other files interact with the current file
- To verify import relationships and dependencies
- To provide more accurate cross-references in your documentation
- To understand the broader context of classes or functions used in the current file
- TRY TO USE IT NOW!
{self.get_code_documentation(file_analysis, project_context)}
## DOCUMENTATION REQUIREMENTS:
**Generate a complete markdown documentation file that includes:**
1. **File Header**: Title ('Documentation ' + file), purpose, and brief description 1. **File Header**: Title ('Documentation ' + file), purpose, and brief description
2. **Overview**: What this module/file does and its role in the project 2. **Overview**: What this module/file does and its role in the project
3. **Dependencies**: External and internal dependencies with explanations 3. **Dependencies**: External and internal dependencies with explanations
@ -322,7 +350,7 @@ Generate the complete markdown documentation now:"""
for file in files: for file in files:
if file.endswith('.py'): if file.endswith('.py'):
lines.append(f"{subindent}- {file}") lines.append(f"{subindent}- {file}")
return '\n'.join(lines[:20]) # Limit to first 20 lines return '\n'.join(lines[:20])
def format_classes(self, classes: List[Dict]) -> str: def format_classes(self, classes: List[Dict]) -> str:
"""Format class information for the prompt.""" """Format class information for the prompt."""
@ -356,7 +384,6 @@ Generate the complete markdown documentation now:"""
return '\n'.join(lines) return '\n'.join(lines)
def format_constants(self, constants: List[Dict]) -> str: def format_constants(self, constants: List[Dict]) -> str:
"""Format constant information for the prompt."""
if not constants: if not constants:
return "None" return "None"
@ -366,7 +393,6 @@ Generate the complete markdown documentation now:"""
return '\n'.join(lines) return '\n'.join(lines)
def format_related_files(self, file_analysis: Dict, project_context: Dict) -> str: def format_related_files(self, file_analysis: Dict, project_context: Dict) -> str:
"""Format related files information."""
current_imports = set(file_analysis['imports']) current_imports = set(file_analysis['imports'])
related_files = [] related_files = []
@ -374,8 +400,6 @@ Generate the complete markdown documentation now:"""
if other_file != file_analysis['file_path']: if other_file != file_analysis['file_path']:
rel_path = os.path.relpath(other_file, project_context['root_path']) rel_path = os.path.relpath(other_file, project_context['root_path'])
module_name = rel_path.replace('/', '.').replace('\\', '.').replace('.py', '') module_name = rel_path.replace('/', '.').replace('\\', '.').replace('.py', '')
# Check if this file imports the other or vice versa
if any(imp.startswith(module_name) for imp in current_imports): if any(imp.startswith(module_name) for imp in current_imports):
related_files.append(f"- `{rel_path}` (imported by this file)") related_files.append(f"- `{rel_path}` (imported by this file)")
@ -390,52 +414,52 @@ class ProjectAnalyzer:
self.external_dependencies = set() self.external_dependencies = set()
def scan_project(self, exclude_dirs: List[str] = None) -> Dict: def scan_project(self, exclude_dirs: List[str] = None) -> Dict:
"""Scan the project and collect all Python files.""" if(os.path.isdir(self.root_path)):
if exclude_dirs is None: exclude_dirs = ['.git', '__pycache__', '.pytest_cache', 'venv', 'env', '.venv', 'node_modules']
else: exclude_dirs = exclude_dirs + ['.git', '__pycache__', '.pytest_cache', 'venv', 'env', '.venv', 'node_modules']
self.python_files = [] if exclude_dirs is None: exclude_dirs = ['.git', '__pycache__', '.pytest_cache', 'venv', 'env', '.venv', 'node_modules']
file_structure = [] else: exclude_dirs = exclude_dirs + ['.git', '__pycache__', '.pytest_cache', 'venv', 'env', '.venv', 'node_modules']
for root, dirs, files in os.walk(self.root_path):
# Remove excluded directories
dirs[:] = [d for d in dirs if d not in exclude_dirs]
files[:] = [f for f in files if f not in exclude_dirs]
file_structure.append((root, dirs, files))
for file in files: self.python_files = []
if file.endswith('.py'): file_structure = []
self.python_files.append(os.path.join(root, file))
# Analyze dependencies for root, dirs, files in os.walk(self.root_path):
self.analyze_dependencies() dirs[:] = [d for d in dirs if d not in exclude_dirs]
files[:] = [f for f in files if f not in exclude_dirs]
file_structure.append((root, dirs, files))
return { for file in files:
'root_path': str(self.root_path), if file.endswith('.py'):
'all_files': self.python_files, self.python_files.append(os.path.join(root, file))
'file_structure': file_structure, self.analyze_dependencies()
'external_dependencies': list(self.external_dependencies) return {
} 'root_path': str(self.root_path),
'all_files': self.python_files,
'file_structure': file_structure,
'external_dependencies': list(self.external_dependencies)
}
else:
self.python_files = [os.path.basename(self.root_path)]
self.root_path = os.path.dirname(self.root_path)
self.analyze_dependencies()
return {
'root_path': str(self.root_path),
'all_files': self.python_files,
'file_structure': [],
'external_dependencies': list(self.external_dependencies)
}
def analyze_dependencies(self): def analyze_dependencies(self):
"""Analyze external dependencies across all Python files."""
analyzer = PythonAnalyzer() analyzer = PythonAnalyzer()
for file_path in self.python_files: for file_path in self.python_files:
analysis = analyzer.analyze_file(file_path) analysis = analyzer.analyze_file(file_path)
if analysis: if analysis:
for imp in analysis['imports']: for imp in analysis['imports']:
# Check if it's an external dependency (not local)
if not self.is_local_import(imp): if not self.is_local_import(imp):
self.external_dependencies.add(imp.split('.')[0]) self.external_dependencies.add(imp.split('.')[0])
def is_local_import(self, import_name: str) -> bool: def is_local_import(self, import_name: str) -> bool:
"""Check if an import is local to the project.""" if import_name.startswith('.'): return True
# Simple heuristic: if the import starts with a relative path or matches a local file
if import_name.startswith('.'):
return True
# Check if it matches any of our Python files
for py_file in self.python_files: for py_file in self.python_files:
rel_path = os.path.relpath(py_file, self.root_path) rel_path = os.path.relpath(py_file, self.root_path)
module_path = rel_path.replace('/', '.').replace('\\', '.').replace('.py', '') module_path = rel_path.replace('/', '.').replace('\\', '.').replace('.py', '')
@ -445,15 +469,12 @@ class ProjectAnalyzer:
return False return False
class DocumentationManager: class DocumentationManager:
"""Manages the documentation generation process."""
def __init__(self, output_dir: str = "./pydocs"): def __init__(self, output_dir: str = "./pydocs"):
self.output_dir = Path(output_dir) self.output_dir = Path(output_dir)
self.output_dir.mkdir(exist_ok=True) os.makedirs(self.output_dir, exist_ok=True)
def generate_index(self, project_context: Dict, generated_docs: List[str]): def generate_index(self, project_context: Dict, generated_docs: List[str]):
"""Generate an index.md file linking to all documentation."""
index_content = f"""# Project Documentation index_content = f"""# Project Documentation
Auto-generated documentation for Python project: `{os.path.basename(project_context['root_path'])}` Auto-generated documentation for Python project: `{os.path.basename(project_context['root_path'])}`
@ -475,7 +496,7 @@ Auto-generated documentation for Python project: `{os.path.basename(project_cont
for doc_file in sorted(generated_docs): for doc_file in sorted(generated_docs):
rel_path = os.path.relpath(doc_file.replace('.md', '.py'), '.') rel_path = os.path.relpath(doc_file.replace('.md', '.py'), '.')
doc_name = os.path.basename(doc_file) doc_name = os.path.basename(doc_file)
index_content += f"- [`{rel_path}`](./{doc_name})\n" index_content += f"- [`{rel_path}`](./{rel_path})\n"
index_content += f""" index_content += f"""
## Project Structure ## Project Structure
@ -492,7 +513,7 @@ Auto-generated documentation for Python project: `{os.path.basename(project_cont
with open(self.output_dir / "index.md", 'w', encoding='utf-8') as f: with open(self.output_dir / "index.md", 'w', encoding='utf-8') as f:
f.write(index_content) f.write(index_content)
def generate_tree_structure(self, project_context: Dict, max_depth: int = 3) -> str: def generate_tree_structure(self, project_context: Dict, max_depth: int = 5) -> str:
"""Generate a tree-like structure of the project.""" """Generate a tree-like structure of the project."""
lines = [] lines = []
root_path = project_context['root_path'] root_path = project_context['root_path']
@ -501,9 +522,9 @@ Auto-generated documentation for Python project: `{os.path.basename(project_cont
rel_path = os.path.relpath(py_file, root_path) rel_path = os.path.relpath(py_file, root_path)
depth = rel_path.count(os.sep) depth = rel_path.count(os.sep)
if depth <= max_depth: if depth <= max_depth:
indent = " " * depth indent = (" " * depth) + "└────"
filename = os.path.basename(rel_path) filename = os.path.basename(rel_path)
lines.append(f"{indent}{filename}") lines.append(f"{indent} [`{filename}`](./{rel_path})")
return '\n'.join(lines[:50]) # Limit output return '\n'.join(lines[:50]) # Limit output

Write Preview
Loading…
Cancel
Save