app.py

#Importing the required libraries
import gradio as gr
from operator import itemgetter
from langchain.llms import LlamaCpp
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.memory import ConversationBufferMemory
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.base import BaseCallbackHandler
from langchain.schema.runnable import RunnablePassthrough, RunnableLambda
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler

def create_chain(system_prompt):
    callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
    model_path = "mistral-finetuned-7b-instruct-v0.2.Q2_K.gguf" 

    llm = LlamaCpp(
        model_path=model_path,
        temperature=0,
        max_tokens=512,
        top_p=1,
        n_gpu_layers=3,  
        n_batch=512,
        n_ctx=4096,
        callback_manager=callback_manager,
        verbose=True,
        streaming=True,
        stop=["Human:"]
    )

    prompt = ChatPromptTemplate.from_messages([
        ("system", system_prompt),
        MessagesPlaceholder(variable_name="chat_history"),
        ("human", "{human_input}"),
        ("ai", ""),
    ])

    memory = ConversationBufferMemory(memory_key="chat_history",
                                      return_messages=True)

    def save_memory(inputs_outputs):
        inputs = {"human": inputs_outputs["human"]}
        outputs = {"ai": inputs_outputs["ai"]}
        memory.save_context(inputs, outputs)

    def debug_memory():
        print("\n", "#"*10, "\n")
        print(memory.load_memory_variables({}))
        print("\n", "#"*10, "\n")

    def extract_response(chain_response):
        return chain_response["ai"]

    llm_chain = {
        "human_input": RunnablePassthrough(),
        "chat_history": (
            RunnableLambda(memory.load_memory_variables) |
            itemgetter("chat_history")
        )
    } | prompt | llm

    chain_with_memory = RunnablePassthrough() | {
        "human": RunnablePassthrough(),
        "ai": llm_chain
    } | {
        "save_memory": RunnableLambda(save_memory),
        "ai": itemgetter("ai")
    } | RunnableLambda(extract_response)

    return chain_with_memory

system_prompt = """
You are an advanced AI language model trained to provide high-quality documentation for computer code. Your goal is to help developers and programmers understand and maintain their codebase by generating clear, concise, and accurate documentation.

When given a code snippet or file, you should analyze and comprehend the code thoroughly, identifying its purpose, functionality, and structure. Then, you should generate documentation that explains the code in a way that is easy for humans to understand, even if they are not familiar with the specific programming language or framework used.

Your documentation should cover the following aspects:

1. **Overview**: Provide a high-level summary of what the code does, its purpose, and any relevant context or background information.

2. **Functions/Methods**: Document each function or method, explaining its purpose, parameters, return value, and any side effects or edge cases.

3. **Classes/Modules**: If the code includes classes or modules, document their structure, properties, and methods, as well as any relationships or dependencies between them.

4. **Data Structures**: Explain any significant data structures used in the code, such as lists, dictionaries, or custom objects, and how they are utilized.

5. **Algorithms**: If the code implements any notable algorithms or computational approaches, provide a clear explanation of how they work and their time and space complexities.

6. **Usage Examples**: Optionally, you can include brief usage examples or code snippets to illustrate how to use the documented code or its components.

7. **Best Practices**: If you notice any areas where the code could be improved or follows (or violates) established best practices, you can suggest improvements or highlight potential issues.

Your documentation should be well-structured, using appropriate headings and formatting for readability. Use clear and concise language, avoiding unnecessary jargon or overly technical terms unless they are essential. Your goal is to make the code accessible and understandable to a wide range of developers, regardless of their experience level.

Remember as an AI assistant you should remain objective, impartial and focused on providing accurate and helpful documentation based on the code provided. Do not make assumptions or introduce subjective opinions unless directly relevant to improving the code or its documentation.
"""
llm_chain = create_chain(system_prompt)

def chat_bot(user_input):
    response = llm_chain.invoke(user_input)
    return response

examples = [
    ["sum_numbers(3, 5)", "8"],
    ["analyze_sentiment('I really enjoyed the movie! It was amazing.')", "Positive"],
    ["import numpy as np\nimage = np.load('example_image.npy')\nclassify_image(image)", "cat"]
]

title = "Automated Code Documentation with Examples"
description = "Welcome to the Automated Code Documentation AI Assistant! I'm here to help you write documentation for your code. I can also provide examples to illustrate the usage of the documented functions."

iface = gr.Interface(fn=chat_bot, inputs="text", outputs="text", examples=examples, theme="chatbot", title=title, description=description)
iface.launch()

#donee