1  """
  2  Documentation specialized agent.
  3  """
  4  
  5  from langchain.agents import AgentExecutor, create_react_agent
  6  from langchain_core.tools import Tool
  7  from langchain_core.prompts import PromptTemplate
  8  from langchain_core.language_models import BaseLLM
  9  
 10  from kamaji.tools import read_file, write_file_smart, list_directory, execute_shell_command
 11  
 12  
 13  def extract_docstrings(file_path: str) -> str:
 14      """Extract all docstrings from a Python file."""
 15      try:
 16          with open(file_path, 'r') as f:
 17              content = f.read()
 18  
 19          import ast
 20          tree = ast.parse(content)
 21  
 22          docstrings = []
 23          for node in ast.walk(tree):
 24              if isinstance(node, (ast.FunctionDef, ast.ClassDef, ast.Module)):
 25                  docstring = ast.get_docstring(node)
 26                  if docstring:
 27                      name = getattr(node, 'name', 'Module')
 28                      docstrings.append(f"{name}: {docstring[:100]}...")
 29  
 30          if docstrings:
 31              return "\n".join(docstrings)
 32          else:
 33              return f"No docstrings found in {file_path}"
 34  
 35      except Exception as e:
 36          return f"Error extracting docstrings: {str(e)}"
 37  
 38  
 39  def check_missing_docs(file_path: str) -> str:
 40      """Find functions/classes without docstrings."""
 41      try:
 42          with open(file_path, 'r') as f:
 43              content = f.read()
 44  
 45          import ast
 46          tree = ast.parse(content)
 47  
 48          missing = []
 49          for node in ast.walk(tree):
 50              if isinstance(node, (ast.FunctionDef, ast.ClassDef)):
 51                  if not ast.get_docstring(node):
 52                      missing.append(f"Line {node.lineno}: {node.name}")
 53  
 54          if missing:
 55              return f"Missing documentation:\n" + "\n".join(missing)
 56          else:
 57              return f"✅ All functions/classes in {file_path} have docstrings"
 58  
 59      except Exception as e:
 60          return f"Error checking documentation: {str(e)}"
 61  
 62  
 63  def create_documentation_agent(llm: BaseLLM, verbose: bool = False) -> AgentExecutor:
 64      """
 65      Create a documentation specialized agent.
 66  
 67      This agent is optimized for writing, reviewing, and improving documentation.
 68  
 69      Args:
 70          llm: The language model to use
 71          verbose: Show detailed reasoning
 72  
 73      Returns:
 74          AgentExecutor configured for documentation tasks
 75      """
 76      tools = [
 77          Tool(
 78              name="read_file",
 79              func=read_file,
 80              description="Read a source file to document. Input should be the file path."
 81          ),
 82          Tool(
 83              name="write_file",
 84              func=write_file_smart,
 85              description="Write documentation to a file. Format: 'filepath|||content'"
 86          ),
 87          Tool(
 88              name="extract_docstrings",
 89              func=extract_docstrings,
 90              description="Extract existing docstrings from a Python file. Input: file path."
 91          ),
 92          Tool(
 93              name="check_missing_docs",
 94              func=check_missing_docs,
 95              description="Find functions/classes without docstrings. Input: file path."
 96          ),
 97          Tool(
 98              name="list_directory",
 99              func=list_directory,
100              description="List files in a directory. Input: directory path."
101          ),
102          Tool(
103              name="run_command",
104              func=execute_shell_command,
105              description="Run documentation tools like sphinx or mkdocs. Input: shell command."
106          ),
107      ]
108  
109      prompt = PromptTemplate.from_template("""
110  You are a technical writer specialized in software documentation.
111  
112  Your responsibilities:
113  - Write clear, comprehensive documentation
114  - Create README files, API docs, and guides
115  - Add/improve docstrings in code
116  - Ensure documentation completeness
117  
118  Documentation standards:
119  📝 Clear and concise language
120  📝 Code examples where helpful
121  📝 Proper formatting (Markdown, reStructuredText)
122  📝 Complete function/class documentation
123  📝 Usage examples and best practices
124  
125  Available tools: {tools}
126  
127  Use the following format:
128  
129  Question: {input}
130  Thought: I should understand what needs to be documented
131  Action: [one of: {tool_names}]
132  Action Input: [input for the tool]
133  Observation: [tool result]
134  ... (repeat Thought/Action/Observation as needed)
135  Thought: I have created/improved the documentation
136  Final Answer: [summary of documentation changes made]
137  
138  Begin!
139  
140  Question: {input}
141  Thought:{agent_scratchpad}
142  """)
143  
144      agent = create_react_agent(llm, tools, prompt)
145  
146      return AgentExecutor(
147          agent=agent,
148          tools=tools,
149          verbose=verbose,
150          handle_parsing_errors=True,
151          max_iterations=30,  # Increased for complex documentation tasks
152          early_stopping_method="generate"  # Generate answer even if not complete
153      )