melghorab/code-assistant

1from langchain_core.tools import Tool, tool2import requests3import json4from docx import Document5import re6from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage7 8 9 10@tool11def search_pypi(package_name: str) -> str:12    """Search PyPI for Python package information. Input should be the package name.13    Args: 14        package_name: name of the package 15    """16    print(f"Tool called for package: {package_name}")17    base_url = "https://pypi.org/pypi"18    try:19        try:20            response = requests.get(f"{base_url}/{package_name}/json")21            response.raise_for_status()22            info = response.json()23        except requests.RequestException as e:24            raise Exception(f"Error fetching PyPI info for {package_name}: {str(e)}")25        result =  json.dumps({26            "name": info["info"]["name"],27            "summary": info["info"]["summary"],28        })29        print(f"Tool result: {result}")30        return result31    except Exception as e:32        return f"Could not find package information: {str(e)}"33    34# @tool35# def write_to_docx(documentation_text: str) -> str:36#     """37#     Writes the AI-generated documentation to a .docx file and returns the file path.38#     """39#     doc = Document()40#     # doc.add_heading("Code Documentation", level=1)41 42#     lines = documentation_text.split("\n")43#     for line in lines:44#         if line.startswith("# "):  # Section Heading45#             doc.add_heading(line[2:], level=1)46#         elif line.startswith("## "):  # Subsection Heading47#             doc.add_heading(line[3:], level=2)48#         else:  # Normal paragraph49#             doc.add_paragraph(line)50 51#     file_path = "generated_documentation.docx"52#     doc.save(file_path)53#     return file_path54 55def write_to_docx(state):56    text = state['messages'][-1].content57    filename = 'generated_documentation.docx'58    doc = Document()59    60    lines = text.split("\n")61    for line in lines:62        if line.startswith("### "):63            doc.add_heading(line[4:], level=3)64        elif line.startswith("## "):65            doc.add_heading(line[3:], level=2)66        elif line.startswith("# "):67            doc.add_heading(line[2:], level=1)68        elif "**" in line:69            bold_parts = re.split(r"(\*\*.*?\*\*)", line)70            para = doc.add_paragraph()71            for part in bold_parts:72                if part.startswith("**") and part.endswith("**"):73                    para.add_run(part[2:-2]).bold = True74                else:75                    para.add_run(part)76        else:77            doc.add_paragraph(line)78    79    # Save document80    doc.save(filename)81    return {"messages": [SystemMessage(content=[filename])]}