model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Implemented using langchain-groq\n",
    "##### Model Used : mixtral-8x7b-32768"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()  # Load environment variables from .env file\n",
    "\n",
    "api_key = os.getenv(\"GROQ_API_KEY\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import PromptTemplate\n",
    "from langchain_groq import ChatGroq\n",
    "from langchain.chains import LLMChain\n",
    "\n",
    "LLM = ChatGroq(groq_api_key=api_key,model=\"mixtral-8x7b-32768\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "must_required_fn=\"\"\"\n",
    "Design Java methods for generating unique user IDs, validating user details, and invoking these two methods to process user input. \n",
    "Details:\n",
    "\n",
    "ID Generation Method: This method should generate unique user IDs by combining the first characters of the user's name, their age, and a static counter value. The counter should increment with each new user to ensure uniqueness.\n",
    "\n",
    "Validation Method: This method should validate user details, ensuring the user is older than 18 and that their phone number is exactly 10 digits long.\n",
    "\n",
    "Invocation Method: This method should call the ID generation and validation methods to process user input data.\n",
    "\n",
    "[INST] Each method must be implemented with appropriate logic and should not be left blank. Also, do not use complex concepts like hash, hashmaps while generating the code [/INST]\n",
    "\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Alternate prompt; currently not in use due to poor results\n",
    "\n",
    "# must_rquired_fn=\"\"\"\n",
    "# Design a Java method that generates unique user IDs by leveraging essential user details and a static counter.  The user ID generation process will involve combining the initial characters of the user's name with their age and appending the current value of a static counter. This counter incrementally updates with each new user instance, ensuring that each generated ID remains distinct and identifiable. Also, you must include a validation method which validates the user based on certain user details like age, phone number, like user must be greater than 18 or phone number must be equal to 10 digits. There must also be a method which invokes the  methods like validation method and id generator method. Also, the program should only take input data from the user. You must implement appropriate logic while defining these methods. Do not leave it blank. Also for better readability, beside the definitions of these functions, add proper comment. You must give comment only to these three functions \n",
    "# \"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Scenario 1 : *Employee Management*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Java Code Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Write a Java program demonstrating inheritance with the following classes:\n",
      "\n",
      "1. Employee (base class): Attributes - name, age, employeeID\n",
      "2. FullTimeEmployee (derived from Employee): Attributes - hourlyRate, hoursWorked\n",
      "3. PartTimeEmployee (derived from Employee): Attributes - salary, department\n",
      "\n",
      "Include appropriate constructors, getters, and setters for each class. In the FullTimeEmployee class, implement a method to calculate the annual salary. In the PartTimeEmployee class, implement a method to calculate the total earnings based on hourlyRate and hoursWorked.\n",
      "\n",
      "Design Java methods for generating unique user IDs, validating user details, and invoking these two methods to process user input. \n",
      "Details:\n",
      "\n",
      "ID Generation Method: This method should generate unique user IDs by combining the first characters of the user's name, their age, and a static counter value. The counter should increment with each new user to ensure uniqueness.\n",
      "\n",
      "Validation Method: This method should validate user details, ensuring the user is older than 18 and that their phone number is exactly 10 digits long.\n",
      "\n",
      "Invocation Method: This method should call the ID generation and validation methods to process user input data.\n",
      "\n",
      "[INST] Each method must be implemented with appropriate logic and should not be left blank. Also, do not use complex concepts like hash, hashmaps while generating the code [/INST]\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#Inheritance \n",
    "\n",
    "prompt_template_inheritance='''\n",
    "Write a Java program demonstrating inheritance with the following classes:\n",
    "\n",
    "1. Employee (base class): Attributes - {Attribute1}, {Attribute2}, {Attribute3}\n",
    "2. FullTimeEmployee (derived from Employee): Attributes - {Attribute4}, {Attribute5}\n",
    "3. PartTimeEmployee (derived from Employee): Attributes - {Attribute6}, {Attribute7}\n",
    "\n",
    "Include appropriate constructors, getters, and setters for each class. In the FullTimeEmployee class, implement a method to calculate the annual salary. In the PartTimeEmployee class, implement a method to calculate the total earnings based on hourlyRate and hoursWorked.\n",
    "'''+must_required_fn\n",
    "\n",
    "inheritance_prompt = PromptTemplate(\n",
    "    input_variables=['Attribute1', 'Attribute2', 'Attribute3', 'Attribute4', 'Attribute5', 'Attribute6', 'Attribute7'],\n",
    "    template=prompt_template_inheritance\n",
    ")\n",
    "\n",
    "formatted_prompt = inheritance_prompt.format(\n",
    "    Attribute1='name',\n",
    "    Attribute2='age',\n",
    "    Attribute3='employeeID',\n",
    "    Attribute4='hourlyRate',\n",
    "    Attribute5='hoursWorked',\n",
    "    Attribute6='salary',\n",
    "    Attribute7='department'\n",
    ")\n",
    "\n",
    "print(formatted_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\Ajish\\Miniconda3\\envs\\prompts@infosys\\Lib\\site-packages\\langchain_core\\_api\\deprecation.py:119: LangChainDeprecationWarning: The class `LLMChain` was deprecated in LangChain 0.1.17 and will be removed in 0.3.0. Use RunnableSequence, e.g., `prompt | llm` instead.\n",
      "  warn_deprecated(\n",
      "c:\\Users\\Ajish\\Miniconda3\\envs\\prompts@infosys\\Lib\\site-packages\\langchain_core\\_api\\deprecation.py:119: LangChainDeprecationWarning: The method `Chain.run` was deprecated in langchain 0.1.0 and will be removed in 0.2.0. Use invoke instead.\n",
      "  warn_deprecated(\n"
     ]
    }
   ],
   "source": [
    "chain = LLMChain(llm=LLM, prompt=inheritance_prompt)\n",
    "raw_output=chain.run({\n",
    "    'Attribute1': 'name',\n",
    "    'Attribute2': 'age',\n",
    "    'Attribute3': 'employeeID',\n",
    "    'Attribute4': 'hourlyRate',\n",
    "    'Attribute5': 'hoursWorked',\n",
    "    'Attribute6': 'salary',\n",
    "    'Attribute7': 'department',\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Here is a simple Java program demonstrating inheritance and the required methods:\n",
      "\n",
      "```java\n",
      "public class Employee {\n",
      "    protected String name;\n",
      "    protected int age;\n",
      "    protected String employeeID;\n",
      "\n",
      "    public Employee(String name, int age) {\n",
      "        this.name = name;\n",
      "        this.age = age;\n",
      "        this.employeeID = generateID(name);\n",
      "    }\n",
      "\n",
      "    public String getName() {\n",
      "        return name;\n",
      "    }\n",
      "\n",
      "    public void setName(String name) {\n",
      "        this.name = name;\n",
      "    }\n",
      "\n",
      "    public int getAge() {\n",
      "        return age;\n",
      "    }\n",
      "\n",
      "    public void setAge(int age) {\n",
      "        this.age = age;\n",
      "    }\n",
      "\n",
      "    public String getEmployeeID() {\n",
      "        return employeeID;\n",
      "    }\n",
      "\n",
      "    private String generateID(String name) {\n",
      "        String counter = Integer.toString(Employee.getIdCounter());\n",
      "        Employee.incrementIdCounter();\n",
      "        return name.charAt(0) + age + counter;\n",
      "    }\n",
      "\n",
      "    private static int idCounter = 1;\n",
      "\n",
      "    public static int getIdCounter() {\n",
      "        return idCounter;\n",
      "    }\n",
      "\n",
      "    public static void incrementIdCounter() {\n",
      "        idCounter++;\n",
      "    }\n",
      "}\n",
      "\n",
      "public class FullTimeEmployee extends Employee {\n",
      "    private double hourlyRate;\n",
      "    private int hoursWorked;\n",
      "\n",
      "    public FullTimeEmployee(String name, int age, double hourlyRate, int hoursWorked) {\n",
      "        super(name, age);\n",
      "        this.hourlyRate = hourlyRate;\n",
      "        this.hoursWorked = hoursWorked;\n",
      "    }\n",
      "\n",
      "    public double getHourlyRate() {\n",
      "        return hourlyRate;\n",
      "    }\n",
      "\n",
      "    public void setHourlyRate(double hourlyRate) {\n",
      "        this.hourlyRate = hourlyRate;\n",
      "    }\n",
      "\n",
      "    public int getHoursWorked() {\n",
      "        return hoursWorked;\n",
      "    }\n",
      "\n",
      "    public void setHoursWorked(int hoursWorked) {\n",
      "        this.hoursWorked = hoursWorked;\n",
      "    }\n",
      "\n",
      "    public double calculateAnnualSalary() {\n",
      "        return hourlyRate * hoursWorked * 52; // Assuming 52 weeks in a year\n",
      "    }\n",
      "}\n",
      "\n",
      "public class PartTimeEmployee extends Employee {\n",
      "    private double salary;\n",
      "    private String department;\n",
      "\n",
      "    public PartTimeEmployee(String name, int age, double salary, String department) {\n",
      "        super(name, age);\n",
      "        this.salary = salary;\n",
      "        this.department = department;\n",
      "    }\n",
      "\n",
      "    public double getSalary() {\n",
      "        return salary;\n",
      "    }\n",
      "\n",
      "    public void setSalary(double salary) {\n",
      "        this.salary = salary;\n",
      "    }\n",
      "\n",
      "    public String getDepartment() {\n",
      "        return department;\n",
      "    }\n",
      "\n",
      "    public void setDepartment(String department) {\n",
      "        this.department = department;\n",
      "    }\n",
      "\n",
      "    public double calculateTotalEarnings() {\n",
      "        // Assuming part-time employees are paid based on their salary only\n",
      "        return salary;\n",
      "    }\n",
      "}\n",
      "\n",
      "public class EmployeeProcessor {\n",
      "    public static boolean validateUserDetails(Employee employee, long phoneNumber) {\n",
      "        if (employee.getAge() < 18) {\n",
      "            System.out.println(\"User is not eligible as they are below 18 years of age.\");\n",
      "            return false;\n",
      "        }\n",
      "        if (String.valueOf(phoneNumber).length() != 10) {\n",
      "            System.out.println(\"User phone number should be exactly 10 digits long.\");\n",
      "            return false;\n",
      "        }\n",
      "        return true;\n",
      "    }\n",
      "\n",
      "    public static void processUserInput(Employee employee, long phoneNumber) {\n",
      "        if (!validateUserDetails((Employee) employee, phoneNumber)) {\n",
      "            return;\n",
      "        }\n",
      "        System.out.println(\"User ID: \" + employee.getEmployeeID());\n",
      "        if (employee instanceof FullTimeEmployee) {\n",
      "            FullTimeEmployee fullTimeEmployee = (FullTimeEmployee) employee;\n",
      "            System.out.println(\"Annual Salary: \" + fullTimeEmployee.calculateAnnualSalary());\n",
      "        } else if (employee instanceof PartTimeEmployee) {\n",
      "            PartTimeEmployee partTimeEmployee = (PartTimeEmployee) employee;\n",
      "            System.out.println(\"Total Earnings: \" + partTimeEmployee.calculateTotalEarnings());\n",
      "        }\n",
      "    }\n",
      "}\n",
      "```\n",
      "\n",
      "This program includes the `Employee`, `FullTimeEmployee`, and `PartTimeEmployee` classes, each with their respective attributes, constructors, getters, and setters. The `EmployeeProcessor` class contains the methods for generating unique user IDs, validating user details, and invoking the calculation methods for full-time and part-time employees.\n",
      "\n",
      "Please note that the validation method does not currently include phone number validation in the format you requested (only 10 digits long). You can add this functionality if needed.\n"
     ]
    }
   ],
   "source": [
    "print(raw_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "public class Employee {\n",
      "    protected String name;\n",
      "    protected int age;\n",
      "    protected String employeeID;\n",
      "\n",
      "    public Employee(String name, int age) {\n",
      "        this.name = name;\n",
      "        this.age = age;\n",
      "        this.employeeID = generateID(name);\n",
      "    }\n",
      "\n",
      "    public String getName() {\n",
      "        return name;\n",
      "    }\n",
      "\n",
      "    public void setName(String name) {\n",
      "        this.name = name;\n",
      "    }\n",
      "\n",
      "    public int getAge() {\n",
      "        return age;\n",
      "    }\n",
      "\n",
      "    public void setAge(int age) {\n",
      "        this.age = age;\n",
      "    }\n",
      "\n",
      "    public String getEmployeeID() {\n",
      "        return employeeID;\n",
      "    }\n",
      "\n",
      "    private String generateID(String name) {\n",
      "        String counter = Integer.toString(Employee.getIdCounter());\n",
      "        Employee.incrementIdCounter();\n",
      "        return name.charAt(0) + age + counter;\n",
      "    }\n",
      "\n",
      "    private static int idCounter = 1;\n",
      "\n",
      "    public static int getIdCounter() {\n",
      "        return idCounter;\n",
      "    }\n",
      "\n",
      "    public static void incrementIdCounter() {\n",
      "        idCounter++;\n",
      "    }\n",
      "}\n",
      "\n",
      "public class FullTimeEmployee extends Employee {\n",
      "    private double hourlyRate;\n",
      "    private int hoursWorked;\n",
      "\n",
      "    public FullTimeEmployee(String name, int age, double hourlyRate, int hoursWorked) {\n",
      "        super(name, age);\n",
      "        this.hourlyRate = hourlyRate;\n",
      "        this.hoursWorked = hoursWorked;\n",
      "    }\n",
      "\n",
      "    public double getHourlyRate() {\n",
      "        return hourlyRate;\n",
      "    }\n",
      "\n",
      "    public void setHourlyRate(double hourlyRate) {\n",
      "        this.hourlyRate = hourlyRate;\n",
      "    }\n",
      "\n",
      "    public int getHoursWorked() {\n",
      "        return hoursWorked;\n",
      "    }\n",
      "\n",
      "    public void setHoursWorked(int hoursWorked) {\n",
      "        this.hoursWorked = hoursWorked;\n",
      "    }\n",
      "\n",
      "    public double calculateAnnualSalary() {\n",
      "        return hourlyRate * hoursWorked * 52; // Assuming 52 weeks in a year\n",
      "    }\n",
      "}\n",
      "\n",
      "public class PartTimeEmployee extends Employee {\n",
      "    private double salary;\n",
      "    private String department;\n",
      "\n",
      "    public PartTimeEmployee(String name, int age, double salary, String department) {\n",
      "        super(name, age);\n",
      "        this.salary = salary;\n",
      "        this.department = department;\n",
      "    }\n",
      "\n",
      "    public double getSalary() {\n",
      "        return salary;\n",
      "    }\n",
      "\n",
      "    public void setSalary(double salary) {\n",
      "        this.salary = salary;\n",
      "    }\n",
      "\n",
      "    public String getDepartment() {\n",
      "        return department;\n",
      "    }\n",
      "\n",
      "    public void setDepartment(String department) {\n",
      "        this.department = department;\n",
      "    }\n",
      "\n",
      "    public double calculateTotalEarnings() {\n",
      "        // Assuming part-time employees are paid based on their salary only\n",
      "        return salary;\n",
      "    }\n",
      "}\n",
      "\n",
      "public class EmployeeProcessor {\n",
      "    public static boolean validateUserDetails(Employee employee, long phoneNumber) {\n",
      "        if (employee.getAge() < 18) {\n",
      "            System.out.println(\"User is not eligible as they are below 18 years of age.\");\n",
      "            return false;\n",
      "        }\n",
      "        if (String.valueOf(phoneNumber).length() != 10) {\n",
      "            System.out.println(\"User phone number should be exactly 10 digits long.\");\n",
      "            return false;\n",
      "        }\n",
      "        return true;\n",
      "    }\n",
      "\n",
      "    public static void processUserInput(Employee employee, long phoneNumber) {\n",
      "        if (!validateUserDetails((Employee) employee, phoneNumber)) {\n",
      "            return;\n",
      "        }\n",
      "        System.out.println(\"User ID: \" + employee.getEmployeeID());\n",
      "        if (employee instanceof FullTimeEmployee) {\n",
      "            FullTimeEmployee fullTimeEmployee = (FullTimeEmployee) employee;\n",
      "            System.out.println(\"Annual Salary: \" + fullTimeEmployee.calculateAnnualSalary());\n",
      "        } else if (employee instanceof PartTimeEmployee) {\n",
      "            PartTimeEmployee partTimeEmployee = (PartTimeEmployee) employee;\n",
      "            System.out.println(\"Total Earnings: \" + partTimeEmployee.calculateTotalEarnings());\n",
      "        }\n",
      "    }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "import re\n",
    "\n",
    "def extract_java_code(text):\n",
    "    \"\"\"\n",
    "    Extracts the Java code block from the provided text.\n",
    "\n",
    "    :param text: str, the input text containing the Java code block\n",
    "    :return: str, the extracted Java code block\n",
    "    \"\"\"\n",
    "    # Use regex to find the Java code block\n",
    "    match = re.search(r'```java(.*?)```', text, re.DOTALL)\n",
    "    if match:\n",
    "        java_code = match.group(1).strip()\n",
    "        return java_code\n",
    "    else:\n",
    "        return \"No Java code block found.\"\n",
    "    \n",
    "prettify_java_code=extract_java_code(raw_output)\n",
    "print(prettify_java_code)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Class Diagram Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generate PlantUML code for creating class diagrams for the following Java program: public class Employee {\n",
      "    protected String name;\n",
      "    protected int age;\n",
      "    protected String employeeID;\n",
      "\n",
      "    public Employee(String name, int age) {\n",
      "        this.name = name;\n",
      "        this.age = age;\n",
      "        this.employeeID = generateID(name);\n",
      "    }\n",
      "\n",
      "    public String getName() {\n",
      "        return name;\n",
      "    }\n",
      "\n",
      "    public void setName(String name) {\n",
      "        this.name = name;\n",
      "    }\n",
      "\n",
      "    public int getAge() {\n",
      "        return age;\n",
      "    }\n",
      "\n",
      "    public void setAge(int age) {\n",
      "        this.age = age;\n",
      "    }\n",
      "\n",
      "    public String getEmployeeID() {\n",
      "        return employeeID;\n",
      "    }\n",
      "\n",
      "    private String generateID(String name) {\n",
      "        String counter = Integer.toString(Employee.getIdCounter());\n",
      "        Employee.incrementIdCounter();\n",
      "        return name.charAt(0) + age + counter;\n",
      "    }\n",
      "\n",
      "    private static int idCounter = 1;\n",
      "\n",
      "    public static int getIdCounter() {\n",
      "        return idCounter;\n",
      "    }\n",
      "\n",
      "    public static void incrementIdCounter() {\n",
      "        idCounter++;\n",
      "    }\n",
      "}\n",
      "\n",
      "public class FullTimeEmployee extends Employee {\n",
      "    private double hourlyRate;\n",
      "    private int hoursWorked;\n",
      "\n",
      "    public FullTimeEmployee(String name, int age, double hourlyRate, int hoursWorked) {\n",
      "        super(name, age);\n",
      "        this.hourlyRate = hourlyRate;\n",
      "        this.hoursWorked = hoursWorked;\n",
      "    }\n",
      "\n",
      "    public double getHourlyRate() {\n",
      "        return hourlyRate;\n",
      "    }\n",
      "\n",
      "    public void setHourlyRate(double hourlyRate) {\n",
      "        this.hourlyRate = hourlyRate;\n",
      "    }\n",
      "\n",
      "    public int getHoursWorked() {\n",
      "        return hoursWorked;\n",
      "    }\n",
      "\n",
      "    public void setHoursWorked(int hoursWorked) {\n",
      "        this.hoursWorked = hoursWorked;\n",
      "    }\n",
      "\n",
      "    public double calculateAnnualSalary() {\n",
      "        return hourlyRate * hoursWorked * 52; // Assuming 52 weeks in a year\n",
      "    }\n",
      "}\n",
      "\n",
      "public class PartTimeEmployee extends Employee {\n",
      "    private double salary;\n",
      "    private String department;\n",
      "\n",
      "    public PartTimeEmployee(String name, int age, double salary, String department) {\n",
      "        super(name, age);\n",
      "        this.salary = salary;\n",
      "        this.department = department;\n",
      "    }\n",
      "\n",
      "    public double getSalary() {\n",
      "        return salary;\n",
      "    }\n",
      "\n",
      "    public void setSalary(double salary) {\n",
      "        this.salary = salary;\n",
      "    }\n",
      "\n",
      "    public String getDepartment() {\n",
      "        return department;\n",
      "    }\n",
      "\n",
      "    public void setDepartment(String department) {\n",
      "        this.department = department;\n",
      "    }\n",
      "\n",
      "    public double calculateTotalEarnings() {\n",
      "        // Assuming part-time employees are paid based on their salary only\n",
      "        return salary;\n",
      "    }\n",
      "}\n",
      "\n",
      "public class EmployeeProcessor {\n",
      "    public static boolean validateUserDetails(Employee employee, long phoneNumber) {\n",
      "        if (employee.getAge() < 18) {\n",
      "            System.out.println(\"User is not eligible as they are below 18 years of age.\");\n",
      "            return false;\n",
      "        }\n",
      "        if (String.valueOf(phoneNumber).length() != 10) {\n",
      "            System.out.println(\"User phone number should be exactly 10 digits long.\");\n",
      "            return false;\n",
      "        }\n",
      "        return true;\n",
      "    }\n",
      "\n",
      "    public static void processUserInput(Employee employee, long phoneNumber) {\n",
      "        if (!validateUserDetails((Employee) employee, phoneNumber)) {\n",
      "            return;\n",
      "        }\n",
      "        System.out.println(\"User ID: \" + employee.getEmployeeID());\n",
      "        if (employee instanceof FullTimeEmployee) {\n",
      "            FullTimeEmployee fullTimeEmployee = (FullTimeEmployee) employee;\n",
      "            System.out.println(\"Annual Salary: \" + fullTimeEmployee.calculateAnnualSalary());\n",
      "        } else if (employee instanceof PartTimeEmployee) {\n",
      "            PartTimeEmployee partTimeEmployee = (PartTimeEmployee) employee;\n",
      "            System.out.println(\"Total Earnings: \" + partTimeEmployee.calculateTotalEarnings());\n",
      "        }\n",
      "    }\n",
      "}. Ensure that the data types of the data members are included in the class diagram. Use the format: `function : itsdatatype` for functions and variables, and for constructors, include the data types of attributes inside the brackets as parameters. Also, don't use any kind of icons, instead use '-' for variables and '+' for functions. Use this format `baseClass  <|-- derivedClass` to represent the relation between two classes.\n"
     ]
    }
   ],
   "source": [
    "prompt_template_classDiagrams = \"Generate PlantUML code for creating class diagrams for the following Java program: {javacode}. Ensure that the data types of the data members are included in the class diagram. Use the format: `function : itsdatatype` for functions and variables, and for constructors, include the data types of attributes inside the brackets as parameters. Also, don't use any kind of icons, instead use '-' for variables and '+' for functions. Use this format `baseClass  <|-- derivedClass` to represent the relation between two classes.\"\n",
    "\n",
    "prompt_classDiagrams = PromptTemplate(\n",
    "    input_variables=['javacode'],\n",
    "    template=prompt_template_classDiagrams\n",
    ")\n",
    "\n",
    "formatted_prompt = prompt_classDiagrams.format(\n",
    "    javacode=prettify_java_code\n",
    ")\n",
    "\n",
    "print(formatted_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = LLMChain(llm=LLM, prompt=prompt_classDiagrams)\n",
    "raw_plantUML_output=chain.run({'javacode':raw_output})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "@startuml\n",
      "\n",
      "class Employee {\n",
      "    - name : String\n",
      "    - age : int\n",
      "    - employeeID : String\n",
      "    + Employee(name : String, age : int)\n",
      "    + getName() : String\n",
      "    + setName(name : String)\n",
      "    + getAge() : int\n",
      "    + setAge(age : int)\n",
      "    + getEmployeeID() : String\n",
      "    + calculateAnnualSalary() : double\n",
      "    + calculateTotalEarnings() : double\n",
      "}\n",
      "\n",
      "class FullTimeEmployee {\n",
      "    - hourlyRate : double\n",
      "    - hoursWorked : int\n",
      "    + FullTimeEmployee(name : String, age : int, hourlyRate : double, hoursWorked : int)\n",
      "    + getHourlyRate() : double\n",
      "    + setHourlyRate(hourlyRate : double)\n",
      "    + getHoursWorked() : int\n",
      "    + setHoursWorked(hoursWorked : int)\n",
      "}\n",
      "\n",
      "class PartTimeEmployee {\n",
      "    - salary : double\n",
      "    - department : String\n",
      "    + PartTimeEmployee(name : String, age : int, salary : double, department : String)\n",
      "    + getSalary() : double\n",
      "    + setSalary(salary : double)\n",
      "    + getDepartment() : String\n",
      "    + setDepartment(department : String)\n",
      "}\n",
      "\n",
      "Employee <|-- FullTimeEmployee\n",
      "Employee <|-- PartTimeEmployee\n",
      "\n",
      "@enduml\n"
     ]
    }
   ],
   "source": [
    "import re\n",
    "\n",
    "def extract_puml_code(text):\n",
    "    # Use regex to find the PlantUML code within triple-backtick fenced code blocks\n",
    "    pattern = r'```plantuml(.*?)```'\n",
    "    match = re.search(pattern, text, re.DOTALL)\n",
    "    if match:\n",
    "        return match.group(1).strip()\n",
    "    else:\n",
    "        return None\n",
    "\n",
    "def prettify_puml_code(puml_code):\n",
    "    # Clean up indentation of PlantUML code for prettification\n",
    "    lines = puml_code.strip().splitlines()\n",
    "    # Determine base indentation level\n",
    "    base_indent = len(lines[0]) - len(lines[0].lstrip())\n",
    "\n",
    "    # Remove base indentation from each line\n",
    "    prettified_lines = [line[base_indent:] for line in lines]\n",
    "\n",
    "    return '\\n'.join(prettified_lines)\n",
    "\n",
    "# Example usage:\n",
    "if __name__ == \"__main__\":\n",
    "    uml = raw_plantUML_output\n",
    "    \n",
    "    puml_code = extract_puml_code(uml)\n",
    "    if puml_code:\n",
    "        prettified_puml = prettify_puml_code(puml_code)\n",
    "        print(prettified_puml)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Class diagram saved as class_diagram.png\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from plantuml import PlantUML, PlantUMLHTTPError\n",
    "\n",
    "def create_class_diagram(plantuml_code, output_file):\n",
    "    \"\"\"\n",
    "    Create a class diagram from PlantUML code and save it as an image.\n",
    "\n",
    "    :param plantuml_code: str, PlantUML code for the class diagram\n",
    "    :param output_file: str, output file path for the generated diagram (e.g., 'diagram.png')\n",
    "    \"\"\"\n",
    "\n",
    "    plantuml_code = \"@startuml\\n\" + \"hide circle\\n\" + \"skinparam classAttributeIconSize 0\\n\" + plantuml_code + \"\\n@enduml\"\n",
    "    \n",
    "    # Save the PlantUML code to a temporary file\n",
    "    temp_file = 'temp_diagram.puml'\n",
    "    with open(temp_file, 'w') as file:\n",
    "        file.write(plantuml_code)\n",
    "    \n",
    "    # Initialize the PlantUML processor (pointing to the PlantUML server)\n",
    "    plantuml = PlantUML(url='http://www.plantuml.com/plantuml/png/')\n",
    "\n",
    "    try:\n",
    "        # Generate the diagram\n",
    "        plantuml.processes_file(temp_file)\n",
    "        \n",
    "        # Define the output image path\n",
    "        output_image = temp_file.replace('.puml', '.png')\n",
    "        \n",
    "        # Move the generated diagram to the desired output file location\n",
    "        if os.path.exists(output_image):\n",
    "            os.rename(output_image, output_file)\n",
    "        else:\n",
    "            raise FileNotFoundError(\"The output image was not generated.\")\n",
    "        \n",
    "        print(f\"Class diagram saved as {output_file}\")\n",
    "\n",
    "    except PlantUMLHTTPError as e:\n",
    "        print(f\"Failed to generate diagram: HTTP error {e.response.status} - {e.response.reason}\")\n",
    "    except Exception as e:\n",
    "        print(f\"An error occurred: {str(e)}\")\n",
    "    finally:\n",
    "        # Clean up the temporary file\n",
    "        if os.path.exists(temp_file):\n",
    "            os.remove(temp_file)\n",
    "\n",
    "# Example usage\n",
    "plantuml_code = prettified_puml\n",
    "\n",
    "output_file = 'class_diagram.png'\n",
    "create_class_diagram(plantuml_code, output_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Document Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_format=\"\"\" \n",
    "<s> [INST] Follow this format for generating the documentation. Do not include any information about getters and setters in the document: [/INST]\n",
    "\n",
    "[INST]\n",
    "**Problem Statement**\n",
    "<Provide a clear and concise problem statement for the program here. Describe what the program is intended to solve or achieve.>\n",
    "\n",
    "**Class Diagram**\n",
    "<Include a class diagram picture from the current directory here (leave this section blank).>\n",
    "\n",
    "**Model Class Documentation**\n",
    "<Repeat the following structure for each class in the program, if applicable.>\n",
    "\n",
    "**Class Name: <class name>**\n",
    "\n",
    "**<name of id generation method>**\n",
    "- This method automatically generates and sets the <name of variable that stores the id>.\n",
    "- <Provide a detailed description of the method, including its purpose, how it works, and any relevant details.>\n",
    "\n",
    "**<name of validation method>**\n",
    "- <Provide a detailed description of the validation method, including what it validates, how it performs the validation, and any specific rules or criteria it follows.>\n",
    "\n",
    "**<name of the method that invokes the validation and id generation>**\n",
    "- <Provide a detailed description of the method that invokes both the validation and ID generation methods. Describe its purpose, how it integrates these methods, and any relevant details about its operation.>\n",
    "\n",
    "[INST]<End of repeat section for classes.>[/INST]\n",
    "[/INST]\n",
    "\n",
    "[INST]Additionally, you must include a detailed description of the validation, ID generation, and the method which invokes both of these methods. Ensure clarity and comprehensiveness in these descriptions.[/INST]\n",
    "\n",
    "[INST]If any of the above mentioned field are not applicable, please skip them and not add any unneccessary information[/INST]\n",
    "\n",
    "[INST]Do not add any unnecessary text or information outside of the specified format.[/INST]\n",
    "</s>\n",
    "\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Alternate prompt; currently not in use due to poor results\n",
    "\n",
    "# doc_format=\"\"\" \n",
    "# <s> [INST] Follow this format for generation of documentation. Also, do not include any information about getters and setters in the document [/INST]:\n",
    "# **Problem Statement**\n",
    "# <problem statement of that program/>\n",
    "# **Class Diagram**\n",
    "# <class diagram picture from the current directory (leave it blank)/>\n",
    "# **Model Class**\n",
    "# <repeat for all classes if applicable>\n",
    "# **<class name>**\n",
    "# **<name of id generation method>**\n",
    "# - This method auto generates and sets the <name of variable which stores the id>.\n",
    "# - <detailed description of the method>\n",
    "\n",
    "# **<name of validation method>**\n",
    "# - <detailed description of the method>\n",
    "\n",
    "# **<name of the method invokes the validation and id generation>**\n",
    "# - <detailed description of the method>\n",
    "# </repeat for all classes if applicable>\n",
    "\n",
    "# [INST]Also, you must include a detailed description about the validation, id generation and the method which invokes both of these methods[/INST]\n",
    "# [INST]Do not add any other texts unneccessarily[/INST]\n",
    "# </s>\n",
    "# \"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " \n",
      "<s> [INST] Follow this format for generating the documentation. Do not include any information about getters and setters in the document: [/INST]\n",
      "\n",
      "[INST]\n",
      "**Problem Statement**\n",
      "<Provide a clear and concise problem statement for the program here. Describe what the program is intended to solve or achieve.>\n",
      "\n",
      "**Class Diagram**\n",
      "<Include a class diagram picture from the current directory here (leave this section blank).>\n",
      "\n",
      "**Model Class Documentation**\n",
      "<Repeat the following structure for each class in the program, if applicable.>\n",
      "\n",
      "**Class Name: <class name>**\n",
      "\n",
      "**<name of id generation method>**\n",
      "- This method automatically generates and sets the <name of variable that stores the id>.\n",
      "- <Provide a detailed description of the method, including its purpose, how it works, and any relevant details.>\n",
      "\n",
      "**<name of validation method>**\n",
      "- <Provide a detailed description of the validation method, including what it validates, how it performs the validation, and any specific rules or criteria it follows.>\n",
      "\n",
      "**<name of the method that invokes the validation and id generation>**\n",
      "- <Provide a detailed description of the method that invokes both the validation and ID generation methods. Describe its purpose, how it integrates these methods, and any relevant details about its operation.>\n",
      "\n",
      "[INST]<End of repeat section for classes.>[/INST]\n",
      "[/INST]\n",
      "\n",
      "[INST]Additionally, you must include a detailed description of the validation, ID generation, and the method which invokes both of these methods. Ensure clarity and comprehensiveness in these descriptions.[/INST]\n",
      "\n",
      "[INST]If any of the above mentioned field are not applicable, please skip them and not add any unneccessary information[/INST]\n",
      "\n",
      "[INST]Do not add any unnecessary text or information outside of the specified format.[/INST]\n",
      "</s>\n",
      "\n",
      "Prepare a documentation for this java program\n",
      "public class Employee {\n",
      "    protected String name;\n",
      "    protected int age;\n",
      "    protected String employeeID;\n",
      "\n",
      "    public Employee(String name, int age) {\n",
      "        this.name = name;\n",
      "        this.age = age;\n",
      "        this.employeeID = generateID(name);\n",
      "    }\n",
      "\n",
      "    public String getName() {\n",
      "        return name;\n",
      "    }\n",
      "\n",
      "    public void setName(String name) {\n",
      "        this.name = name;\n",
      "    }\n",
      "\n",
      "    public int getAge() {\n",
      "        return age;\n",
      "    }\n",
      "\n",
      "    public void setAge(int age) {\n",
      "        this.age = age;\n",
      "    }\n",
      "\n",
      "    public String getEmployeeID() {\n",
      "        return employeeID;\n",
      "    }\n",
      "\n",
      "    private String generateID(String name) {\n",
      "        String counter = Integer.toString(Employee.getIdCounter());\n",
      "        Employee.incrementIdCounter();\n",
      "        return name.charAt(0) + age + counter;\n",
      "    }\n",
      "\n",
      "    private static int idCounter = 1;\n",
      "\n",
      "    public static int getIdCounter() {\n",
      "        return idCounter;\n",
      "    }\n",
      "\n",
      "    public static void incrementIdCounter() {\n",
      "        idCounter++;\n",
      "    }\n",
      "}\n",
      "\n",
      "public class FullTimeEmployee extends Employee {\n",
      "    private double hourlyRate;\n",
      "    private int hoursWorked;\n",
      "\n",
      "    public FullTimeEmployee(String name, int age, double hourlyRate, int hoursWorked) {\n",
      "        super(name, age);\n",
      "        this.hourlyRate = hourlyRate;\n",
      "        this.hoursWorked = hoursWorked;\n",
      "    }\n",
      "\n",
      "    public double getHourlyRate() {\n",
      "        return hourlyRate;\n",
      "    }\n",
      "\n",
      "    public void setHourlyRate(double hourlyRate) {\n",
      "        this.hourlyRate = hourlyRate;\n",
      "    }\n",
      "\n",
      "    public int getHoursWorked() {\n",
      "        return hoursWorked;\n",
      "    }\n",
      "\n",
      "    public void setHoursWorked(int hoursWorked) {\n",
      "        this.hoursWorked = hoursWorked;\n",
      "    }\n",
      "\n",
      "    public double calculateAnnualSalary() {\n",
      "        return hourlyRate * hoursWorked * 52; // Assuming 52 weeks in a year\n",
      "    }\n",
      "}\n",
      "\n",
      "public class PartTimeEmployee extends Employee {\n",
      "    private double salary;\n",
      "    private String department;\n",
      "\n",
      "    public PartTimeEmployee(String name, int age, double salary, String department) {\n",
      "        super(name, age);\n",
      "        this.salary = salary;\n",
      "        this.department = department;\n",
      "    }\n",
      "\n",
      "    public double getSalary() {\n",
      "        return salary;\n",
      "    }\n",
      "\n",
      "    public void setSalary(double salary) {\n",
      "        this.salary = salary;\n",
      "    }\n",
      "\n",
      "    public String getDepartment() {\n",
      "        return department;\n",
      "    }\n",
      "\n",
      "    public void setDepartment(String department) {\n",
      "        this.department = department;\n",
      "    }\n",
      "\n",
      "    public double calculateTotalEarnings() {\n",
      "        // Assuming part-time employees are paid based on their salary only\n",
      "        return salary;\n",
      "    }\n",
      "}\n",
      "\n",
      "public class EmployeeProcessor {\n",
      "    public static boolean validateUserDetails(Employee employee, long phoneNumber) {\n",
      "        if (employee.getAge() < 18) {\n",
      "            System.out.println(\"User is not eligible as they are below 18 years of age.\");\n",
      "            return false;\n",
      "        }\n",
      "        if (String.valueOf(phoneNumber).length() != 10) {\n",
      "            System.out.println(\"User phone number should be exactly 10 digits long.\");\n",
      "            return false;\n",
      "        }\n",
      "        return true;\n",
      "    }\n",
      "\n",
      "    public static void processUserInput(Employee employee, long phoneNumber) {\n",
      "        if (!validateUserDetails((Employee) employee, phoneNumber)) {\n",
      "            return;\n",
      "        }\n",
      "        System.out.println(\"User ID: \" + employee.getEmployeeID());\n",
      "        if (employee instanceof FullTimeEmployee) {\n",
      "            FullTimeEmployee fullTimeEmployee = (FullTimeEmployee) employee;\n",
      "            System.out.println(\"Annual Salary: \" + fullTimeEmployee.calculateAnnualSalary());\n",
      "        } else if (employee instanceof PartTimeEmployee) {\n",
      "            PartTimeEmployee partTimeEmployee = (PartTimeEmployee) employee;\n",
      "            System.out.println(\"Total Earnings: \" + partTimeEmployee.calculateTotalEarnings());\n",
      "        }\n",
      "    }\n",
      "} .\n"
     ]
    }
   ],
   "source": [
    "doc_prompt_template=doc_format+\"\"\"Prepare a documentation for this java program\n",
    "{javacode} .\"\"\"\n",
    "\n",
    "\n",
    "prompt_documentation = PromptTemplate(\n",
    "    input_variables=['javacode'],\n",
    "    template=doc_prompt_template\n",
    ")\n",
    "\n",
    "formatted_prompt = prompt_documentation.format(\n",
    "    javacode=prettify_java_code\n",
    ")\n",
    "\n",
    "print(formatted_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = LLMChain(llm=LLM, prompt=prompt_documentation)\n",
    "raw_docs=chain.run({'javacode':prettify_java_code})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prepare a documentation for this java program\n",
      "\n",
      "**Problem Statement**\n",
      "\n",
      "The program is intended to manage a list of employees and perform validation checks on their age and phone number, calculate their earnings based on whether they are full-time or part-time, and generate a unique ID for each employee.\n",
      "\n",
      "**Class Diagram**\n",
      "\n",
      "<Insert class diagram picture here>\n",
      "\n",
      "**Model Class Documentation**\n",
      "\n",
      "**Class Name: Employee**\n",
      "\n",
      "**employeeID**\n",
      "- This variable stores the unique ID of the employee.\n",
      "\n",
      "**name**\n",
      "- This variable stores the name of the employee.\n",
      "\n",
      "**age**\n",
      "- This variable stores the age of the employee.\n",
      "\n",
      "**generateID(String name)**\n",
      "- This method generates a unique ID for the employee using the first character of the employee's name, their age, and a counter that is incremented each time a new ID is generated.\n",
      "- The generated ID is assigned to the employeeID variable.\n",
      "\n",
      "**getIdCounter()**\n",
      "- This method returns the current value of the counter used for generating unique IDs.\n",
      "\n",
      "**incrementIdCounter()**\n",
      "- This method increments the counter used for generating unique IDs.\n",
      "\n",
      "**Class Name: FullTimeEmployee**\n",
      "\n",
      "**hourlyRate**\n",
      "- This variable stores the hourly rate of the full-time employee.\n",
      "\n",
      "**hoursWorked**\n",
      "- This variable stores the number of hours worked by the full-time employee.\n",
      "\n",
      "**calculateAnnualSalary()**\n",
      "- This method calculates the annual salary of the full-time employee by multiplying their hourly rate by the number of hours worked and the number of weeks in a year.\n",
      "\n",
      "**Class Name: PartTimeEmployee**\n",
      "\n",
      "**salary**\n",
      "- This variable stores the salary of the part-time employee.\n",
      "\n",
      "**department**\n",
      "- This variable stores the department of the part-time employee.\n",
      "\n",
      "**calculateTotalEarnings()**\n",
      "- This method calculates the total earnings of the part-time employee based on their salary.\n",
      "\n",
      "**Class Name: EmployeeProcessor**\n",
      "\n",
      "**validateUserDetails(Employee employee, long phoneNumber)**\n",
      "- This method validates the age and phone number of the employee.\n",
      "- It checks if the employee's age is less than 18, and if the phone number is not exactly 10 digits long.\n",
      "- If the validation fails, it prints an error message and returns false.\n",
      "\n",
      "**processUserInput(Employee employee, long phoneNumber)**\n",
      "- This method invokes the validateUserDetails method to perform validation checks on the employee's age and phone number.\n",
      "- If the validation fails, it returns.\n",
      "- If the validation passes, it prints the employee's ID and calculates their earnings based on their type, whether full-time or part-time.\n"
     ]
    }
   ],
   "source": [
    "print(raw_docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document saved as Java_Program_Documentation.docx\n"
     ]
    }
   ],
   "source": [
    "from docx import Document\n",
    "from docx.shared import RGBColor\n",
    "from docx.enum.text import WD_PARAGRAPH_ALIGNMENT\n",
    "\n",
    "def get_page_width(doc):\n",
    "    # Get the page width based on left and right margins\n",
    "    return doc.sections[0].page_width - doc.sections[0].left_margin - doc.sections[0].right_margin\n",
    "\n",
    "def separate_formatting_rules(doc_string):\n",
    "    parts = doc_string.split('**Formatting Rules**')\n",
    "    main_text = parts[0].strip()\n",
    "    formatting_rules = parts[1].strip() if len(parts) > 1 else \"\"\n",
    "    return main_text, formatting_rules\n",
    "\n",
    "def create_java_program_doc(doc_string, class_diagram_path, file_name='Java_Program_Documentation.docx'):\n",
    "    # Separate the formatting rules from the main text\n",
    "    main_text, formatting_rules = separate_formatting_rules(doc_string)\n",
    "\n",
    "    # Create a new Document\n",
    "    doc = Document()\n",
    "\n",
    "    # Split the main text into lines\n",
    "    lines = main_text.split('\\n')\n",
    "    \n",
    "    for line in lines:\n",
    "        if line.startswith('**Problem Statement**'):\n",
    "            # Problem statement in bold\n",
    "            paragraph = doc.add_paragraph()\n",
    "            run = paragraph.add_run(line.strip(\"**\"))\n",
    "            run.bold = True\n",
    "            run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "        elif line.startswith('**Class Diagram**'):\n",
    "            # Class diagram section\n",
    "            paragraph = doc.add_paragraph()\n",
    "            run = paragraph.add_run(line.strip(\"**\"))\n",
    "            run.bold = True\n",
    "            run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "            # Add class diagram image\n",
    "            page_width = get_page_width(doc)\n",
    "            doc.add_picture(class_diagram_path, width=page_width)\n",
    "        elif line.startswith('**Model Class**'):\n",
    "            # Model class section\n",
    "            paragraph = doc.add_paragraph()\n",
    "            run = paragraph.add_run(line.strip(\"**\"))\n",
    "            run.bold = True\n",
    "            run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "        elif line.startswith('**') and line.endswith('**'):\n",
    "            # Bold class names\n",
    "            paragraph = doc.add_paragraph()\n",
    "            run = paragraph.add_run(line.strip(\"**\"))\n",
    "            run.bold = True\n",
    "            run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "        elif line.startswith('* '):\n",
    "            # Bulleted list item\n",
    "            paragraph = doc.add_paragraph(line.strip(\"* \"), style='List Bullet')\n",
    "            for run in paragraph.runs:\n",
    "                run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "        elif line.startswith('    * '):\n",
    "            # Nested bullet list item\n",
    "            paragraph = doc.add_paragraph(line.strip().replace(\"*\",\"\"), style='List Bullet 2')\n",
    "            for run in paragraph.runs:\n",
    "                run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "        else:\n",
    "            # Regular paragraph\n",
    "            paragraph = doc.add_paragraph(line)\n",
    "            for run in paragraph.runs:\n",
    "                run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "\n",
    "    # Save the document\n",
    "    doc.save(file_name)\n",
    "    print(f\"Document saved as {file_name}\")\n",
    "\n",
    "# Example usage\n",
    "doc_string = raw_docs\n",
    "class_diagram_path = \"class_diagram.png\"  # Path to your class diagram PNG file\n",
    "create_java_program_doc(doc_string, class_diagram_path, file_name='Java_Program_Documentation.docx')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Scenario 2 : *Order Management*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Java Code Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Write a Java program for an Order class, specifying attributes such as orderID, customerName, totalAmount, paymentMethod, and status. Ensure the inclusion of methods for generating a unique ID and validating an order based on the following conditions:\n",
      "1. orderID should be exactly 12 characters long.\n",
      "2. customerName should not be empty.\n",
      "3. totalAmount should be a positive number.\n",
      "4. paymentMethod should be one of the predefined payment methods (e.g., \"Credit Card\", \"PayPal\", \"Cash\").\n",
      "5. status should be one of the predefined statuses (e.g., \"Pending\", \"Confirmed\", \"Shipped\").\n",
      "\n",
      "Design Java methods for generating unique user IDs, validating user details, and invoking these two methods to process user input. \n",
      "Details:\n",
      "\n",
      "ID Generation Method: This method should generate unique user IDs by combining the first characters of the user's name, their age, and a static counter value. The counter should increment with each new user to ensure uniqueness.\n",
      "\n",
      "Validation Method: This method should validate user details, ensuring the user is older than 18 and that their phone number is exactly 10 digits long.\n",
      "\n",
      "Invocation Method: This method should call the ID generation and validation methods to process user input data.\n",
      "\n",
      "[INST] Each method must be implemented with appropriate logic and should not be left blank. Also, do not use complex concepts like hash, hashmaps while generating the code [/INST]\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#Order Details\n",
    "\n",
    "# Define the prompt template for the Order class\n",
    "prompt_template_order = \"\"\"\n",
    "Write a Java program for an {classname} class, specifying attributes such as {Attribute1}, {Attribute2}, {Attribute3}, {Attribute4}, and {Attribute5}. Ensure the inclusion of methods for generating a unique ID and validating an order based on the following conditions:\n",
    "1. {Attribute1} should be exactly 12 characters long.\n",
    "2. {Attribute2} should not be empty.\n",
    "3. {Attribute3} should be a positive number.\n",
    "4. {Attribute4} should be one of the predefined payment methods (e.g., \"Credit Card\", \"PayPal\", \"Cash\").\n",
    "5. {Attribute5} should be one of the predefined statuses (e.g., \"Pending\", \"Confirmed\", \"Shipped\").\n",
    "\"\"\"+must_required_fn\n",
    "\n",
    "# Define the PromptTemplate instance for the Order class\n",
    "order_prompt = PromptTemplate(\n",
    "    input_variables=['classname', 'Attribute1', 'Attribute2', 'Attribute3', 'Attribute4', 'Attribute5'],\n",
    "    template=prompt_template_order\n",
    ")\n",
    "\n",
    "# Format the prompt with specific attributes for the Order class\n",
    "formatted_prompt = order_prompt.format(\n",
    "    classname='Order',\n",
    "    Attribute1='orderID',\n",
    "    Attribute2='customerName',\n",
    "    Attribute3='totalAmount',\n",
    "    Attribute4='paymentMethod',\n",
    "    Attribute5='status'\n",
    ")\n",
    "\n",
    "print(formatted_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = LLMChain(llm=LLM, prompt=order_prompt)\n",
    "raw_output=chain.run({\n",
    "    'classname':'Order',\n",
    "    'Attribute1':'orderID',\n",
    "    'Attribute2':'customerName',\n",
    "    'Attribute3':'totalAmount',\n",
    "    'Attribute4':'paymentMethod',\n",
    "    'Attribute5':'status'\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "public class Order {\n",
      "    private static int counter = 1;\n",
      "    private String orderID;\n",
      "    private String customerName;\n",
      "    private double totalAmount;\n",
      "    private String paymentMethod;\n",
      "    private String status;\n",
      "\n",
      "    public Order() {\n",
      "        generateOrderID();\n",
      "        this.customerName = \"\";\n",
      "        this.totalAmount = 0.0;\n",
      "        this.paymentMethod = \"\";\n",
      "        this.status = \"\";\n",
      "    }\n",
      "\n",
      "    public void generateOrderID() {\n",
      "        String id = \"\";\n",
      "        // Assuming name is not empty and age is available from some other class or input\n",
      "        String name = \"John Doe\"; // replace with actual name\n",
      "        int age = 25; // replace with actual age\n",
      "\n",
      "        id += name.substring(0, 1); // first character of name\n",
      "        id += age; // age\n",
      "\n",
      "        id += counter++; // static counter\n",
      "\n",
      "        if (id.length() != 12) {\n",
      "            // Pad with zeros if length is less than 12\n",
      "            id = String.format(\"%12s\", id);\n",
      "        }\n",
      "\n",
      "        this.orderID = id;\n",
      "    }\n",
      "\n",
      "    public boolean validateOrder(String name, double amount, String payment, String status) {\n",
      "        if (name.isEmpty() || amount <= 0 || payment.equals(\"\") || status.equals(\"\")) {\n",
      "            return false;\n",
      "        }\n",
      "\n",
      "        String[] acceptedPaymentMethods = {\"Credit Card\", \"PayPal\", \"Cash\"};\n",
      "        String[] acceptedStatuses = {\"Pending\", \"Confirmed\", \"Shipped\"};\n",
      "\n",
      "        for (String method : acceptedPaymentMethods) {\n",
      "            if (payment.equals(method)) {\n",
      "                this.paymentMethod = payment;\n",
      "                break;\n",
      "            }\n",
      "        }\n",
      "\n",
      "        for (String stat : acceptedStatuses) {\n",
      "            if (status.equals(stat)) {\n",
      "                this.status = status;\n",
      "                break;\n",
      "            }\n",
      "        }\n",
      "\n",
      "        this.customerName = name;\n",
      "        this.totalAmount = amount;\n",
      "\n",
      "        return true;\n",
      "    }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "import re\n",
    "\n",
    "def extract_java_code(text):\n",
    "    \"\"\"\n",
    "    Extracts the Java code block from the provided text.\n",
    "\n",
    "    :param text: str, the input text containing the Java code block\n",
    "    :return: str, the extracted Java code block\n",
    "    \"\"\"\n",
    "    # Use regex to find the Java code block\n",
    "    match = re.search(r'```java(.*?)```', text, re.DOTALL)\n",
    "    if match:\n",
    "        java_code = match.group(1).strip()\n",
    "        return java_code\n",
    "    else:\n",
    "        return \"No Java code block found.\"\n",
    "    \n",
    "prettify_java_code=extract_java_code(raw_output)\n",
    "print(prettify_java_code)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Class Diagram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generate PlantUML code for creating class diagrams for the following Java program: public class Order {\n",
      "    private static int counter = 1;\n",
      "    private String orderID;\n",
      "    private String customerName;\n",
      "    private double totalAmount;\n",
      "    private String paymentMethod;\n",
      "    private String status;\n",
      "\n",
      "    public Order() {\n",
      "        generateOrderID();\n",
      "        this.customerName = \"\";\n",
      "        this.totalAmount = 0.0;\n",
      "        this.paymentMethod = \"\";\n",
      "        this.status = \"\";\n",
      "    }\n",
      "\n",
      "    public void generateOrderID() {\n",
      "        String id = \"\";\n",
      "        // Assuming name is not empty and age is available from some other class or input\n",
      "        String name = \"John Doe\"; // replace with actual name\n",
      "        int age = 25; // replace with actual age\n",
      "\n",
      "        id += name.substring(0, 1); // first character of name\n",
      "        id += age; // age\n",
      "\n",
      "        id += counter++; // static counter\n",
      "\n",
      "        if (id.length() != 12) {\n",
      "            // Pad with zeros if length is less than 12\n",
      "            id = String.format(\"%12s\", id);\n",
      "        }\n",
      "\n",
      "        this.orderID = id;\n",
      "    }\n",
      "\n",
      "    public boolean validateOrder(String name, double amount, String payment, String status) {\n",
      "        if (name.isEmpty() || amount <= 0 || payment.equals(\"\") || status.equals(\"\")) {\n",
      "            return false;\n",
      "        }\n",
      "\n",
      "        String[] acceptedPaymentMethods = {\"Credit Card\", \"PayPal\", \"Cash\"};\n",
      "        String[] acceptedStatuses = {\"Pending\", \"Confirmed\", \"Shipped\"};\n",
      "\n",
      "        for (String method : acceptedPaymentMethods) {\n",
      "            if (payment.equals(method)) {\n",
      "                this.paymentMethod = payment;\n",
      "                break;\n",
      "            }\n",
      "        }\n",
      "\n",
      "        for (String stat : acceptedStatuses) {\n",
      "            if (status.equals(stat)) {\n",
      "                this.status = status;\n",
      "                break;\n",
      "            }\n",
      "        }\n",
      "\n",
      "        this.customerName = name;\n",
      "        this.totalAmount = amount;\n",
      "\n",
      "        return true;\n",
      "    }\n",
      "}. Ensure that the data types of the data members are included in the class diagram. Use the format: `function : itsdatatype` for functions and variables, and for constructors, include the data types of attributes inside the brackets as parameters. Also, don't use any kind of icons, instead use '-' for variables and '+' for functions. Use this format `baseClass  <|-- derivedClass` to represent the relation between two classes.\n"
     ]
    }
   ],
   "source": [
    "prompt_template_classDiagrams = \"Generate PlantUML code for creating class diagrams for the following Java program: {javacode}. Ensure that the data types of the data members are included in the class diagram. Use the format: `function : itsdatatype` for functions and variables, and for constructors, include the data types of attributes inside the brackets as parameters. Also, don't use any kind of icons, instead use '-' for variables and '+' for functions. Use this format `baseClass  <|-- derivedClass` to represent the relation between two classes.\"\n",
    "\n",
    "prompt_classDiagrams = PromptTemplate(\n",
    "    input_variables=['javacode'],\n",
    "    template=prompt_template_classDiagrams\n",
    ")\n",
    "\n",
    "formatted_prompt = prompt_classDiagrams.format(\n",
    "    javacode=prettify_java_code\n",
    ")\n",
    "\n",
    "print(formatted_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = LLMChain(llm=LLM, prompt=prompt_classDiagrams)\n",
    "raw_plantUML_output=chain.run({'javacode':raw_output})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "@startuml\n",
      "\n",
      "class Order {\n",
      "    -orderID : String\n",
      "    -customerName : String\n",
      "    -totalAmount : double\n",
      "    -paymentMethod : String\n",
      "    -status : String\n",
      "    +Order()\n",
      "    +generateOrderID()\n",
      "    +validateOrder(name: String, amount: double, payment: String, status: String) : boolean\n",
      "}\n",
      "\n",
      "@enduml\n"
     ]
    }
   ],
   "source": [
    "import re\n",
    "\n",
    "def extract_puml_code(text):\n",
    "    # Use regex to find the PlantUML code within triple-backtick fenced code blocks\n",
    "    pattern = r'```plantuml(.*?)```'\n",
    "    match = re.search(pattern, text, re.DOTALL)\n",
    "    if match:\n",
    "        return match.group(1).strip()\n",
    "    else:\n",
    "        return None\n",
    "\n",
    "def prettify_puml_code(puml_code):\n",
    "    # Clean up indentation of PlantUML code for prettification\n",
    "    lines = puml_code.strip().splitlines()\n",
    "    # Determine base indentation level\n",
    "    base_indent = len(lines[0]) - len(lines[0].lstrip())\n",
    "\n",
    "    # Remove base indentation from each line\n",
    "    prettified_lines = [line[base_indent:] for line in lines]\n",
    "\n",
    "    return '\\n'.join(prettified_lines)\n",
    "\n",
    "# Example usage:\n",
    "if __name__ == \"__main__\":\n",
    "    uml = raw_plantUML_output\n",
    "    \n",
    "    puml_code = extract_puml_code(uml)\n",
    "    if puml_code:\n",
    "        prettified_puml = prettify_puml_code(puml_code)\n",
    "        print(prettified_puml)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Class diagram saved as class_diagram.png\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from plantuml import PlantUML, PlantUMLHTTPError\n",
    "\n",
    "def create_class_diagram(plantuml_code, output_file):\n",
    "    \"\"\"\n",
    "    Create a class diagram from PlantUML code and save it as an image.\n",
    "\n",
    "    :param plantuml_code: str, PlantUML code for the class diagram\n",
    "    :param output_file: str, output file path for the generated diagram (e.g., 'diagram.png')\n",
    "    \"\"\"\n",
    "\n",
    "    plantuml_code = \"@startuml\\n\" + \"hide circle\\n\" + \"skinparam classAttributeIconSize 0\\n\" + plantuml_code + \"\\n@enduml\"\n",
    "    \n",
    "    # Save the PlantUML code to a temporary file\n",
    "    temp_file = 'temp_diagram.puml'\n",
    "    with open(temp_file, 'w') as file:\n",
    "        file.write(plantuml_code)\n",
    "    \n",
    "    # Initialize the PlantUML processor (pointing to the PlantUML server)\n",
    "    plantuml = PlantUML(url='http://www.plantuml.com/plantuml/png/')\n",
    "\n",
    "    try:\n",
    "        # Generate the diagram\n",
    "        plantuml.processes_file(temp_file)\n",
    "        \n",
    "        # Define the output image path\n",
    "        output_image = temp_file.replace('.puml', '.png')\n",
    "        \n",
    "        # Move the generated diagram to the desired output file location\n",
    "        if os.path.exists(output_image):\n",
    "            os.rename(output_image, output_file)\n",
    "        else:\n",
    "            raise FileNotFoundError(\"The output image was not generated.\")\n",
    "        \n",
    "        print(f\"Class diagram saved as {output_file}\")\n",
    "\n",
    "    except PlantUMLHTTPError as e:\n",
    "        print(f\"Failed to generate diagram: HTTP error {e.response.status} - {e.response.reason}\")\n",
    "    except Exception as e:\n",
    "        print(f\"An error occurred: {str(e)}\")\n",
    "    finally:\n",
    "        # Clean up the temporary file\n",
    "        if os.path.exists(temp_file):\n",
    "            os.remove(temp_file)\n",
    "\n",
    "# Example usage\n",
    "plantuml_code = prettified_puml\n",
    "\n",
    "output_file = 'class_diagram.png'\n",
    "create_class_diagram(plantuml_code, output_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Documentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_format=\"\"\" \n",
    "<s> [INST] Follow this format for generating the documentation. Do not include any information about getters and setters in the document: [/INST]\n",
    "\n",
    "[INST]\n",
    "**Problem Statement**\n",
    "<Provide a clear and concise problem statement for the program here. Describe what the program is intended to solve or achieve.>\n",
    "\n",
    "**Class Diagram**\n",
    "<Include a class diagram picture from the current directory here (leave this section blank).>\n",
    "\n",
    "**Model Class Documentation**\n",
    "<Repeat the following structure for each class in the program, if applicable.>\n",
    "\n",
    "**Class Name: <class name>**\n",
    "\n",
    "**<name of id generation method>**\n",
    "- This method automatically generates and sets the <name of variable that stores the id>.\n",
    "- <Provide a detailed description of the method, including its purpose, how it works, and any relevant details.>\n",
    "\n",
    "**<name of validation method>**\n",
    "- <Provide a detailed description of the validation method, including what it validates, how it performs the validation, and any specific rules or criteria it follows.>\n",
    "\n",
    "**<name of the method that invokes the validation and id generation>**\n",
    "- <Provide a detailed description of the method that invokes both the validation and ID generation methods. Describe its purpose, how it integrates these methods, and any relevant details about its operation.>\n",
    "\n",
    "<End of repeat section for classes.>\n",
    "[/INST]\n",
    "\n",
    "[INST]Additionally, you must include a detailed description of the validation, ID generation, and the method which invokes both of these methods. Ensure clarity and comprehensiveness in these descriptions.[/INST]\n",
    "\n",
    "[INST]If any of the above mentioned field are not applicable, please skip them and not add any unneccessary information[/INST]\n",
    "\n",
    "[INST]Do not add any unnecessary text or information outside of the specified format.[/INST]\n",
    "</s>\n",
    "\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prepare a documentation for this java program\n",
      "public class Order {\n",
      "    private static int counter = 1;\n",
      "    private String orderID;\n",
      "    private String customerName;\n",
      "    private double totalAmount;\n",
      "    private String paymentMethod;\n",
      "    private String status;\n",
      "\n",
      "    public Order() {\n",
      "        generateOrderID();\n",
      "        this.customerName = \"\";\n",
      "        this.totalAmount = 0.0;\n",
      "        this.paymentMethod = \"\";\n",
      "        this.status = \"\";\n",
      "    }\n",
      "\n",
      "    public void generateOrderID() {\n",
      "        String id = \"\";\n",
      "        // Assuming name is not empty and age is available from some other class or input\n",
      "        String name = \"John Doe\"; // replace with actual name\n",
      "        int age = 25; // replace with actual age\n",
      "\n",
      "        id += name.substring(0, 1); // first character of name\n",
      "        id += age; // age\n",
      "\n",
      "        id += counter++; // static counter\n",
      "\n",
      "        if (id.length() != 12) {\n",
      "            // Pad with zeros if length is less than 12\n",
      "            id = String.format(\"%12s\", id);\n",
      "        }\n",
      "\n",
      "        this.orderID = id;\n",
      "    }\n",
      "\n",
      "    public boolean validateOrder(String name, double amount, String payment, String status) {\n",
      "        if (name.isEmpty() || amount <= 0 || payment.equals(\"\") || status.equals(\"\")) {\n",
      "            return false;\n",
      "        }\n",
      "\n",
      "        String[] acceptedPaymentMethods = {\"Credit Card\", \"PayPal\", \"Cash\"};\n",
      "        String[] acceptedStatuses = {\"Pending\", \"Confirmed\", \"Shipped\"};\n",
      "\n",
      "        for (String method : acceptedPaymentMethods) {\n",
      "            if (payment.equals(method)) {\n",
      "                this.paymentMethod = payment;\n",
      "                break;\n",
      "            }\n",
      "        }\n",
      "\n",
      "        for (String stat : acceptedStatuses) {\n",
      "            if (status.equals(stat)) {\n",
      "                this.status = status;\n",
      "                break;\n",
      "            }\n",
      "        }\n",
      "\n",
      "        this.customerName = name;\n",
      "        this.totalAmount = amount;\n",
      "\n",
      "        return true;\n",
      "    }\n",
      "} . \n",
      "<s> [INST] Follow this format for generating the documentation. Do not include any information about getters and setters in the document: [/INST]\n",
      "\n",
      "[INST]\n",
      "**Problem Statement**\n",
      "<Provide a clear and concise problem statement for the program here. Describe what the program is intended to solve or achieve.>\n",
      "\n",
      "**Class Diagram**\n",
      "<Include a class diagram picture from the current directory here (leave this section blank).>\n",
      "\n",
      "**Model Class Documentation**\n",
      "<Repeat the following structure for each class in the program, if applicable.>\n",
      "\n",
      "**Class Name: <class name>**\n",
      "\n",
      "**<name of id generation method>**\n",
      "- This method automatically generates and sets the <name of variable that stores the id>.\n",
      "- <Provide a detailed description of the method, including its purpose, how it works, and any relevant details.>\n",
      "\n",
      "**<name of validation method>**\n",
      "- <Provide a detailed description of the validation method, including what it validates, how it performs the validation, and any specific rules or criteria it follows.>\n",
      "\n",
      "**<name of the method that invokes the validation and id generation>**\n",
      "- <Provide a detailed description of the method that invokes both the validation and ID generation methods. Describe its purpose, how it integrates these methods, and any relevant details about its operation.>\n",
      "\n",
      "<End of repeat section for classes.>\n",
      "[/INST]\n",
      "\n",
      "[INST]Additionally, you must include a detailed description of the validation, ID generation, and the method which invokes both of these methods. Ensure clarity and comprehensiveness in these descriptions.[/INST]\n",
      "\n",
      "[INST]If any of the above mentioned field are not applicable, please skip them and not add any unneccessary information[/INST]\n",
      "\n",
      "[INST]Do not add any unnecessary text or information outside of the specified format.[/INST]\n",
      "</s>\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "doc_prompt_template=\"\"\"Prepare a documentation for this java program\n",
    "{javacode} .\"\"\"+doc_format\n",
    "\n",
    "prompt_documentation = PromptTemplate(\n",
    "    input_variables=['javacode'],\n",
    "    template=doc_prompt_template\n",
    ")\n",
    "\n",
    "formatted_prompt = prompt_documentation.format(\n",
    "    javacode=prettify_java_code\n",
    ")\n",
    "\n",
    "print(formatted_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = LLMChain(llm=LLM, prompt=prompt_documentation)\n",
    "raw_docs=chain.run({'javacode':prettify_java_code})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document saved as Java_Program_Documentation.docx\n"
     ]
    }
   ],
   "source": [
    "from docx import Document\n",
    "from docx.shared import RGBColor\n",
    "from docx.enum.text import WD_PARAGRAPH_ALIGNMENT\n",
    "\n",
    "def get_page_width(doc):\n",
    "    # Get the page width based on left and right margins\n",
    "    return doc.sections[0].page_width - doc.sections[0].left_margin - doc.sections[0].right_margin\n",
    "\n",
    "def separate_formatting_rules(doc_string):\n",
    "    parts = doc_string.split('**Formatting Rules**')\n",
    "    main_text = parts[0].strip()\n",
    "    formatting_rules = parts[1].strip() if len(parts) > 1 else \"\"\n",
    "    return main_text, formatting_rules\n",
    "\n",
    "def create_java_program_doc(doc_string, class_diagram_path, file_name='Java_Program_Documentation.docx'):\n",
    "    # Separate the formatting rules from the main text\n",
    "    main_text, formatting_rules = separate_formatting_rules(doc_string)\n",
    "\n",
    "    # Create a new Document\n",
    "    doc = Document()\n",
    "\n",
    "    # Split the main text into lines\n",
    "    lines = main_text.split('\\n')\n",
    "    \n",
    "    for line in lines:\n",
    "        if line.startswith('**Problem Statement**'):\n",
    "            # Problem statement in bold\n",
    "            paragraph = doc.add_paragraph()\n",
    "            run = paragraph.add_run(line.strip(\"**\"))\n",
    "            run.bold = True\n",
    "            run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "        elif line.startswith('**Class Diagram**'):\n",
    "            # Class diagram section\n",
    "            paragraph = doc.add_paragraph()\n",
    "            run = paragraph.add_run(line.strip(\"**\"))\n",
    "            run.bold = True\n",
    "            run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "            # Add class diagram image\n",
    "            page_width = get_page_width(doc)\n",
    "            doc.add_picture(class_diagram_path, width=page_width)\n",
    "        elif line.startswith('**Model Class**'):\n",
    "            # Model class section\n",
    "            paragraph = doc.add_paragraph()\n",
    "            run = paragraph.add_run(line.strip(\"**\"))\n",
    "            run.bold = True\n",
    "            run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "        elif line.startswith('**') and line.endswith('**'):\n",
    "            # Bold class names\n",
    "            paragraph = doc.add_paragraph()\n",
    "            run = paragraph.add_run(line.strip(\"**\"))\n",
    "            run.bold = True\n",
    "            run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "        elif line.startswith('* '):\n",
    "            # Bulleted list item\n",
    "            paragraph = doc.add_paragraph(line.strip(\"* \"), style='List Bullet')\n",
    "            for run in paragraph.runs:\n",
    "                run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "        elif line.startswith('    * '):\n",
    "            # Nested bullet list item\n",
    "            paragraph = doc.add_paragraph(line.strip().replace(\"*\",\"\"), style='List Bullet 2')\n",
    "            for run in paragraph.runs:\n",
    "                run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "        else:\n",
    "            # Regular paragraph\n",
    "            paragraph = doc.add_paragraph(line)\n",
    "            for run in paragraph.runs:\n",
    "                run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "\n",
    "    # Save the document\n",
    "    doc.save(file_name)\n",
    "    print(f\"Document saved as {file_name}\")\n",
    "\n",
    "# Example usage\n",
    "doc_string = raw_docs\n",
    "class_diagram_path = \"class_diagram.png\"  # Path to your class diagram PNG file\n",
    "create_java_program_doc(doc_string, class_diagram_path, file_name='Java_Program_Documentation.docx')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Scenario 3:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Java Code Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 137,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Write a Java program demonstrating inheritance with the following classes:\n",
      "\n",
      "1. Customer (base class): Attributes - name, age, customerID\n",
      "2. RegularCustomer (derived from Customer): Attributes - purchaseAmount, discountRate\n",
      "3. PremiumCustomer (derived from Customer): Attributes - purchaseAmount, loyaltyPoints\n",
      "\n",
      "Include appropriate constructors, getters, and setters for each class. In the RegularCustomer class, implement a method to calculate the discount based on purchase amount. In the PremiumCustomer class, implement a method to calculate the loyalty points based on purchase amount.\n",
      "\n",
      "Design Java methods for generating unique user IDs, validating user details, and invoking these two methods to process user input. \n",
      "Details:\n",
      "\n",
      "ID Generation Method: This method should generate unique user IDs by combining the first characters of the user's name, their age, and a static counter value. The counter should increment with each new user to ensure uniqueness.\n",
      "\n",
      "Validation Method: This method should validate user details, ensuring the user is older than 18 and that their phone number is exactly 10 digits long.\n",
      "\n",
      "Invocation Method: This method should call the ID generation and validation methods to process user input data.\n",
      "\n",
      "[INST] Each method must be implemented with appropriate logic and should not be left blank. Also, do not use complex concepts like hash, hashmaps while generating the code [/INST]\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#Customer\n",
    "prompt_template_inheritance = '''\n",
    "Write a Java program demonstrating inheritance with the following classes:\n",
    "\n",
    "1. Customer (base class): Attributes - {Attribute1}, {Attribute2}, {Attribute3}\n",
    "2. RegularCustomer (derived from Customer): Attributes - {Attribute4}, {Attribute5}\n",
    "3. PremiumCustomer (derived from Customer): Attributes - {Attribute6}, {Attribute7}\n",
    "\n",
    "Include appropriate constructors, getters, and setters for each class. In the RegularCustomer class, implement a method to calculate the discount based on purchase amount. In the PremiumCustomer class, implement a method to calculate the loyalty points based on purchase amount.\n",
    "''' + must_required_fn\n",
    "\n",
    "inheritance_prompt = PromptTemplate(\n",
    "    input_variables=['Attribute1', 'Attribute2', 'Attribute3', 'Attribute4', 'Attribute5', 'Attribute6', 'Attribute7'],\n",
    "    template=prompt_template_inheritance\n",
    ")\n",
    "\n",
    "formatted_prompt = inheritance_prompt.format(\n",
    "    Attribute1='name',\n",
    "    Attribute2='age',\n",
    "    Attribute3='customerID',\n",
    "    Attribute4='purchaseAmount',\n",
    "    Attribute5='discountRate',\n",
    "    Attribute6='purchaseAmount',\n",
    "    Attribute7='loyaltyPoints'\n",
    ")\n",
    "print(formatted_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the LLMChain and invoke it\n",
    "chain = LLMChain(llm=LLM, prompt=inheritance_prompt)\n",
    "raw_output = chain.run({\n",
    "    'Attribute1': 'name',\n",
    "    'Attribute2': 'age',\n",
    "    'Attribute3': 'customerID',\n",
    "    'Attribute4': 'purchaseAmount',\n",
    "    'Attribute5': 'discountRate',\n",
    "    'Attribute6': 'purchaseAmount',\n",
    "    'Attribute7': 'loyaltyPoints',\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "public class Customer {\n",
      "    protected String name;\n",
      "    protected int age;\n",
      "    protected String customerID;\n",
      "\n",
      "    public Customer(String name, int age) {\n",
      "        this.name = name;\n",
      "        this.age = age;\n",
      "        this.customerID = generateCustomerID();\n",
      "    }\n",
      "\n",
      "    public String getName() {\n",
      "        return name;\n",
      "    }\n",
      "\n",
      "    public void setName(String name) {\n",
      "        this.name = name;\n",
      "    }\n",
      "\n",
      "    public int getAge() {\n",
      "        return age;\n",
      "    }\n",
      "\n",
      "    public void setAge(int age) {\n",
      "        this.age = age;\n",
      "    }\n",
      "\n",
      "    public String getCustomerID() {\n",
      "        return customerID;\n",
      "    }\n",
      "\n",
      "    private String generateCustomerID() {\n",
      "        String firstChar = name.substring(0, 1);\n",
      "        return firstChar + age + Counter.counter;\n",
      "    }\n",
      "}\n",
      "\n",
      "class RegularCustomer extends Customer {\n",
      "    private double purchaseAmount;\n",
      "    private final double DISCOUNT_THRESHOLD = 100;\n",
      "    private final double DISCOUNT_RATE = 0.05;\n",
      "\n",
      "    public RegularCustomer(String name, int age, double purchaseAmount) {\n",
      "        super(name, age);\n",
      "        this.purchaseAmount = purchaseAmount;\n",
      "    }\n",
      "\n",
      "    public double getPurchaseAmount() {\n",
      "        return purchaseAmount;\n",
      "    }\n",
      "\n",
      "    public void setPurchaseAmount(double purchaseAmount) {\n",
      "        this.purchaseAmount = purchaseAmount;\n",
      "    }\n",
      "\n",
      "    public double getDiscount() {\n",
      "        if (purchaseAmount > DISCOUNT_THRESHOLD) {\n",
      "            return DISCOUNT_RATE * purchaseAmount;\n",
      "        } else {\n",
      "            return 0;\n",
      "        }\n",
      "    }\n",
      "}\n",
      "\n",
      "class PremiumCustomer extends Customer {\n",
      "    private double purchaseAmount;\n",
      "    private final double LOYALTY_POINT_RATE = 0.02;\n",
      "\n",
      "    public PremiumCustomer(String name, int age, double purchaseAmount) {\n",
      "        super(name, age);\n",
      "        this.purchaseAmount = purchaseAmount;\n",
      "    }\n",
      "\n",
      "    public double getPurchaseAmount() {\n",
      "        return purchaseAmount;\n",
      "    }\n",
      "\n",
      "    public void setPurchaseAmount(double purchaseAmount) {\n",
      "        this.purchaseAmount = purchaseAmount;\n",
      "    }\n",
      "\n",
      "    public int getLoyaltyPoints() {\n",
      "        return (int) (LOYALTY_POINT_RATE * purchaseAmount);\n",
      "    }\n",
      "}\n",
      "\n",
      "class Counter {\n",
      "    public static int counter = 0;\n",
      "}\n",
      "\n",
      "public class Main {\n",
      "    public static void main(String[] args) {\n",
      "        RegularCustomer rc = new RegularCustomer(\"John Doe\", 25, 200);\n",
      "        System.out.println(\"Regular Customer Name: \" + rc.getName());\n",
      "        System.out.println(\"Regular Customer Age: \" + rc.getAge());\n",
      "        System.out.println(\"Regular Customer ID: \" + rc.getCustomerID());\n",
      "        System.out.println(\"Regular Customer Purchase Amount: \" + rc.getPurchaseAmount());\n",
      "        System.out.println(\"Regular Customer Discount: \" + rc.getDiscount());\n",
      "\n",
      "        PremiumCustomer pc = new PremiumCustomer(\"Jane Doe\", 30, 300);\n",
      "        System.out.println(\"\\nPremium Customer Name: \" + pc.getName());\n",
      "        System.out.println(\"Premium Customer Age: \" + pc.getAge());\n",
      "        System.out.println(\"Premium Customer ID: \" + pc.getCustomerID());\n",
      "        System.out.println(\"Premium Customer Purchase Amount: \" + pc.getPurchaseAmount());\n",
      "        System.out.println(\"Premium Customer Loyalty Points: \" + pc.getLoyaltyPoints());\n",
      "    }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "import re\n",
    "\n",
    "def extract_java_code(text):\n",
    "    \"\"\"\n",
    "    Extracts the Java code block from the provided text.\n",
    "\n",
    "    :param text: str, the input text containing the Java code block\n",
    "    :return: str, the extracted Java code block\n",
    "    \"\"\"\n",
    "    # Use regex to find the Java code block\n",
    "    match = re.search(r'```java(.*?)```', text, re.DOTALL)\n",
    "    if match:\n",
    "        java_code = match.group(1).strip()\n",
    "        return java_code\n",
    "    else:\n",
    "        return \"No Java code block found.\"\n",
    "    \n",
    "prettify_java_code=extract_java_code(raw_output)\n",
    "print(prettify_java_code)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Class Diagram Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generate PlantUML code for creating class diagrams for the following Java program: public class Customer {\n",
      "    protected String name;\n",
      "    protected int age;\n",
      "    protected String customerID;\n",
      "\n",
      "    public Customer(String name, int age) {\n",
      "        this.name = name;\n",
      "        this.age = age;\n",
      "        this.customerID = generateCustomerID();\n",
      "    }\n",
      "\n",
      "    public String getName() {\n",
      "        return name;\n",
      "    }\n",
      "\n",
      "    public void setName(String name) {\n",
      "        this.name = name;\n",
      "    }\n",
      "\n",
      "    public int getAge() {\n",
      "        return age;\n",
      "    }\n",
      "\n",
      "    public void setAge(int age) {\n",
      "        this.age = age;\n",
      "    }\n",
      "\n",
      "    public String getCustomerID() {\n",
      "        return customerID;\n",
      "    }\n",
      "\n",
      "    private String generateCustomerID() {\n",
      "        String firstChar = name.substring(0, 1);\n",
      "        return firstChar + age + Counter.counter;\n",
      "    }\n",
      "}\n",
      "\n",
      "class RegularCustomer extends Customer {\n",
      "    private double purchaseAmount;\n",
      "    private final double DISCOUNT_THRESHOLD = 100;\n",
      "    private final double DISCOUNT_RATE = 0.05;\n",
      "\n",
      "    public RegularCustomer(String name, int age, double purchaseAmount) {\n",
      "        super(name, age);\n",
      "        this.purchaseAmount = purchaseAmount;\n",
      "    }\n",
      "\n",
      "    public double getPurchaseAmount() {\n",
      "        return purchaseAmount;\n",
      "    }\n",
      "\n",
      "    public void setPurchaseAmount(double purchaseAmount) {\n",
      "        this.purchaseAmount = purchaseAmount;\n",
      "    }\n",
      "\n",
      "    public double getDiscount() {\n",
      "        if (purchaseAmount > DISCOUNT_THRESHOLD) {\n",
      "            return DISCOUNT_RATE * purchaseAmount;\n",
      "        } else {\n",
      "            return 0;\n",
      "        }\n",
      "    }\n",
      "}\n",
      "\n",
      "class PremiumCustomer extends Customer {\n",
      "    private double purchaseAmount;\n",
      "    private final double LOYALTY_POINT_RATE = 0.02;\n",
      "\n",
      "    public PremiumCustomer(String name, int age, double purchaseAmount) {\n",
      "        super(name, age);\n",
      "        this.purchaseAmount = purchaseAmount;\n",
      "    }\n",
      "\n",
      "    public double getPurchaseAmount() {\n",
      "        return purchaseAmount;\n",
      "    }\n",
      "\n",
      "    public void setPurchaseAmount(double purchaseAmount) {\n",
      "        this.purchaseAmount = purchaseAmount;\n",
      "    }\n",
      "\n",
      "    public int getLoyaltyPoints() {\n",
      "        return (int) (LOYALTY_POINT_RATE * purchaseAmount);\n",
      "    }\n",
      "}\n",
      "\n",
      "class Counter {\n",
      "    public static int counter = 0;\n",
      "}\n",
      "\n",
      "public class Main {\n",
      "    public static void main(String[] args) {\n",
      "        RegularCustomer rc = new RegularCustomer(\"John Doe\", 25, 200);\n",
      "        System.out.println(\"Regular Customer Name: \" + rc.getName());\n",
      "        System.out.println(\"Regular Customer Age: \" + rc.getAge());\n",
      "        System.out.println(\"Regular Customer ID: \" + rc.getCustomerID());\n",
      "        System.out.println(\"Regular Customer Purchase Amount: \" + rc.getPurchaseAmount());\n",
      "        System.out.println(\"Regular Customer Discount: \" + rc.getDiscount());\n",
      "\n",
      "        PremiumCustomer pc = new PremiumCustomer(\"Jane Doe\", 30, 300);\n",
      "        System.out.println(\"\\nPremium Customer Name: \" + pc.getName());\n",
      "        System.out.println(\"Premium Customer Age: \" + pc.getAge());\n",
      "        System.out.println(\"Premium Customer ID: \" + pc.getCustomerID());\n",
      "        System.out.println(\"Premium Customer Purchase Amount: \" + pc.getPurchaseAmount());\n",
      "        System.out.println(\"Premium Customer Loyalty Points: \" + pc.getLoyaltyPoints());\n",
      "    }\n",
      "}. Ensure that the data types of the data members are included in the class diagram. Use the format: `function : itsdatatype` for functions and variables, and for constructors, include the data types of attributes inside the brackets as parameters. Also, don't use any kind of icons, instead use '-' for variables and '+' for functions. Use this format `baseClass  <|-- derivedClass` to represent the relation between two classes.\n"
     ]
    }
   ],
   "source": [
    "prompt_template_classDiagrams = \"Generate PlantUML code for creating class diagrams for the following Java program: {javacode}. Ensure that the data types of the data members are included in the class diagram. Use the format: `function : itsdatatype` for functions and variables, and for constructors, include the data types of attributes inside the brackets as parameters. Also, don't use any kind of icons, instead use '-' for variables and '+' for functions. Use this format `baseClass  <|-- derivedClass` to represent the relation between two classes.\"\n",
    "\n",
    "prompt_classDiagrams = PromptTemplate(\n",
    "    input_variables=['javacode'],\n",
    "    template=prompt_template_classDiagrams\n",
    ")\n",
    "\n",
    "formatted_prompt = prompt_classDiagrams.format(\n",
    "    javacode=prettify_java_code\n",
    ")\n",
    "\n",
    "print(formatted_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = LLMChain(llm=LLM, prompt=prompt_classDiagrams)\n",
    "raw_plantUML_output=chain.run({'javacode':raw_output})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "@startuml\n",
      "\n",
      "class Customer {\n",
      "    - name : String\n",
      "    - age : int\n",
      "    - customerID : String\n",
      "    + Customer(name : String, age : int)\n",
      "    + getName() : String\n",
      "    + setName(name : String)\n",
      "    + getAge() : int\n",
      "    + setAge(age : int)\n",
      "    + getCustomerID() : String\n",
      "}\n",
      "\n",
      "class RegularCustomer {\n",
      "    - purchaseAmount : double\n",
      "    + RegularCustomer(name : String, age : int, purchaseAmount : double)\n",
      "    + getPurchaseAmount() : double\n",
      "    + setPurchaseAmount(purchaseAmount : double)\n",
      "    + getDiscount() : double\n",
      "}\n",
      "\n",
      "class PremiumCustomer {\n",
      "    - purchaseAmount : double\n",
      "    + PremiumCustomer(name : String, age : int, purchaseAmount : double)\n",
      "    + getPurchaseAmount() : double\n",
      "    + setPurchaseAmount(purchaseAmount : double)\n",
      "    + getLoyaltyPoints() : int\n",
      "}\n",
      "\n",
      "class Counter {\n",
      "    + counter : int\n",
      "}\n",
      "\n",
      "Customer <|-- RegularCustomer\n",
      "Customer <|-- PremiumCustomer\n",
      "\n",
      "@enduml\n"
     ]
    }
   ],
   "source": [
    "import re\n",
    "\n",
    "def extract_puml_code(text):\n",
    "    # Use regex to find the PlantUML code within triple-backtick fenced code blocks\n",
    "    pattern = r'```plantuml(.*?)```'\n",
    "    match = re.search(pattern, text, re.DOTALL)\n",
    "    if match:\n",
    "        return match.group(1).strip()\n",
    "    else:\n",
    "        return None\n",
    "\n",
    "def prettify_puml_code(puml_code):\n",
    "    # Clean up indentation of PlantUML code for prettification\n",
    "    lines = puml_code.strip().splitlines()\n",
    "    # Determine base indentation level\n",
    "    base_indent = len(lines[0]) - len(lines[0].lstrip())\n",
    "\n",
    "    # Remove base indentation from each line\n",
    "    prettified_lines = [line[base_indent:] for line in lines]\n",
    "\n",
    "    return '\\n'.join(prettified_lines)\n",
    "\n",
    "# Example usage:\n",
    "if __name__ == \"__main__\":\n",
    "    uml = raw_plantUML_output\n",
    "    \n",
    "    puml_code = extract_puml_code(uml)\n",
    "    if puml_code:\n",
    "        prettified_puml = prettify_puml_code(puml_code)\n",
    "        print(prettified_puml)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Class diagram saved as class_diagram.png\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from plantuml import PlantUML, PlantUMLHTTPError\n",
    "\n",
    "def create_class_diagram(plantuml_code, output_file):\n",
    "    \"\"\"\n",
    "    Create a class diagram from PlantUML code and save it as an image.\n",
    "\n",
    "    :param plantuml_code: str, PlantUML code for the class diagram\n",
    "    :param output_file: str, output file path for the generated diagram (e.g., 'diagram.png')\n",
    "    \"\"\"\n",
    "\n",
    "    plantuml_code = \"@startuml\\n\" + \"hide circle\\n\" + \"skinparam classAttributeIconSize 0\\n\" + plantuml_code + \"\\n@enduml\"\n",
    "    \n",
    "    # Save the PlantUML code to a temporary file\n",
    "    temp_file = 'temp_diagram.puml'\n",
    "    with open(temp_file, 'w') as file:\n",
    "        file.write(plantuml_code)\n",
    "    \n",
    "    # Initialize the PlantUML processor (pointing to the PlantUML server)\n",
    "    plantuml = PlantUML(url='http://www.plantuml.com/plantuml/png/')\n",
    "\n",
    "    try:\n",
    "        # Generate the diagram\n",
    "        plantuml.processes_file(temp_file)\n",
    "        \n",
    "        # Define the output image path\n",
    "        output_image = temp_file.replace('.puml', '.png')\n",
    "        \n",
    "        # Move the generated diagram to the desired output file location\n",
    "        if os.path.exists(output_image):\n",
    "            os.rename(output_image, output_file)\n",
    "        else:\n",
    "            raise FileNotFoundError(\"The output image was not generated.\")\n",
    "        \n",
    "        print(f\"Class diagram saved as {output_file}\")\n",
    "\n",
    "    except PlantUMLHTTPError as e:\n",
    "        print(f\"Failed to generate diagram: HTTP error {e.response.status} - {e.response.reason}\")\n",
    "    except Exception as e:\n",
    "        print(f\"An error occurred: {str(e)}\")\n",
    "    finally:\n",
    "        # Clean up the temporary file\n",
    "        if os.path.exists(temp_file):\n",
    "            os.remove(temp_file)\n",
    "\n",
    "# Example usage\n",
    "plantuml_code = prettified_puml\n",
    "\n",
    "output_file = 'class_diagram.png'\n",
    "create_class_diagram(plantuml_code, output_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Document Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_format=\"\"\" \n",
    "<s> [INST] Follow this format for generating the documentation. Do not include any information about getters and setters in the document: [/INST]\n",
    "\n",
    "[INST]\n",
    "**Problem Statement**\n",
    "<Provide a clear and concise problem statement for the program here. Describe what the program is intended to solve or achieve.>\n",
    "\n",
    "**Class Diagram**\n",
    "<Include a class diagram picture from the current directory here (leave this section blank).>\n",
    "\n",
    "**Model Class Documentation**\n",
    "<Repeat the following structure for each class in the program, if applicable.>\n",
    "\n",
    "**Class Name: <class name>**\n",
    "\n",
    "**<name of id generation method>**\n",
    "- This method automatically generates and sets the <name of variable that stores the id>.\n",
    "- <Provide a detailed description of the method, including its purpose, how it works, and any relevant details.>\n",
    "\n",
    "**<name of validation method>**\n",
    "- <Provide a detailed description of the validation method, including what it validates, how it performs the validation, and any specific rules or criteria it follows.>\n",
    "\n",
    "**<name of the method that invokes the validation and id generation>**\n",
    "- <Provide a detailed description of the method that invokes both the validation and ID generation methods. Describe its purpose, how it integrates these methods, and any relevant details about its operation.>\n",
    "\n",
    "<End of repeat section for classes.>\n",
    "[/INST]\n",
    "\n",
    "[INST]Additionally, you must include a detailed description of the validation, ID generation, and the method which invokes both of these methods. Ensure clarity and comprehensiveness in these descriptions.[/INST]\n",
    "\n",
    "[INST]If any of the above mentioned field are not applicable, please skip them and not add any unneccessary information[/INST]\n",
    "\n",
    "[INST]Do not add any unnecessary text or information outside of the specified format.[/INST]\n",
    "</s>\n",
    "\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prepare a documentation for this java program\n",
      "public class Customer {\n",
      "    protected String name;\n",
      "    protected int age;\n",
      "    protected String customerID;\n",
      "\n",
      "    public Customer(String name, int age) {\n",
      "        this.name = name;\n",
      "        this.age = age;\n",
      "        this.customerID = generateCustomerID();\n",
      "    }\n",
      "\n",
      "    public String getName() {\n",
      "        return name;\n",
      "    }\n",
      "\n",
      "    public void setName(String name) {\n",
      "        this.name = name;\n",
      "    }\n",
      "\n",
      "    public int getAge() {\n",
      "        return age;\n",
      "    }\n",
      "\n",
      "    public void setAge(int age) {\n",
      "        this.age = age;\n",
      "    }\n",
      "\n",
      "    public String getCustomerID() {\n",
      "        return customerID;\n",
      "    }\n",
      "\n",
      "    private String generateCustomerID() {\n",
      "        String firstChar = name.substring(0, 1);\n",
      "        return firstChar + age + Counter.counter;\n",
      "    }\n",
      "}\n",
      "\n",
      "class RegularCustomer extends Customer {\n",
      "    private double purchaseAmount;\n",
      "    private final double DISCOUNT_THRESHOLD = 100;\n",
      "    private final double DISCOUNT_RATE = 0.05;\n",
      "\n",
      "    public RegularCustomer(String name, int age, double purchaseAmount) {\n",
      "        super(name, age);\n",
      "        this.purchaseAmount = purchaseAmount;\n",
      "    }\n",
      "\n",
      "    public double getPurchaseAmount() {\n",
      "        return purchaseAmount;\n",
      "    }\n",
      "\n",
      "    public void setPurchaseAmount(double purchaseAmount) {\n",
      "        this.purchaseAmount = purchaseAmount;\n",
      "    }\n",
      "\n",
      "    public double getDiscount() {\n",
      "        if (purchaseAmount > DISCOUNT_THRESHOLD) {\n",
      "            return DISCOUNT_RATE * purchaseAmount;\n",
      "        } else {\n",
      "            return 0;\n",
      "        }\n",
      "    }\n",
      "}\n",
      "\n",
      "class PremiumCustomer extends Customer {\n",
      "    private double purchaseAmount;\n",
      "    private final double LOYALTY_POINT_RATE = 0.02;\n",
      "\n",
      "    public PremiumCustomer(String name, int age, double purchaseAmount) {\n",
      "        super(name, age);\n",
      "        this.purchaseAmount = purchaseAmount;\n",
      "    }\n",
      "\n",
      "    public double getPurchaseAmount() {\n",
      "        return purchaseAmount;\n",
      "    }\n",
      "\n",
      "    public void setPurchaseAmount(double purchaseAmount) {\n",
      "        this.purchaseAmount = purchaseAmount;\n",
      "    }\n",
      "\n",
      "    public int getLoyaltyPoints() {\n",
      "        return (int) (LOYALTY_POINT_RATE * purchaseAmount);\n",
      "    }\n",
      "}\n",
      "\n",
      "class Counter {\n",
      "    public static int counter = 0;\n",
      "}\n",
      "\n",
      "public class Main {\n",
      "    public static void main(String[] args) {\n",
      "        RegularCustomer rc = new RegularCustomer(\"John Doe\", 25, 200);\n",
      "        System.out.println(\"Regular Customer Name: \" + rc.getName());\n",
      "        System.out.println(\"Regular Customer Age: \" + rc.getAge());\n",
      "        System.out.println(\"Regular Customer ID: \" + rc.getCustomerID());\n",
      "        System.out.println(\"Regular Customer Purchase Amount: \" + rc.getPurchaseAmount());\n",
      "        System.out.println(\"Regular Customer Discount: \" + rc.getDiscount());\n",
      "\n",
      "        PremiumCustomer pc = new PremiumCustomer(\"Jane Doe\", 30, 300);\n",
      "        System.out.println(\"\\nPremium Customer Name: \" + pc.getName());\n",
      "        System.out.println(\"Premium Customer Age: \" + pc.getAge());\n",
      "        System.out.println(\"Premium Customer ID: \" + pc.getCustomerID());\n",
      "        System.out.println(\"Premium Customer Purchase Amount: \" + pc.getPurchaseAmount());\n",
      "        System.out.println(\"Premium Customer Loyalty Points: \" + pc.getLoyaltyPoints());\n",
      "    }\n",
      "} . \n",
      "<s> [INST] Follow this format for generating the documentation. Do not include any information about getters and setters in the document: [/INST]\n",
      "\n",
      "[INST]\n",
      "**Problem Statement**\n",
      "<Provide a clear and concise problem statement for the program here. Describe what the program is intended to solve or achieve.>\n",
      "\n",
      "**Class Diagram**\n",
      "<Include a class diagram picture from the current directory here (leave this section blank).>\n",
      "\n",
      "**Model Class Documentation**\n",
      "<Repeat the following structure for each class in the program, if applicable.>\n",
      "\n",
      "**Class Name: <class name>**\n",
      "\n",
      "**<name of id generation method>**\n",
      "- This method automatically generates and sets the <name of variable that stores the id>.\n",
      "- <Provide a detailed description of the method, including its purpose, how it works, and any relevant details.>\n",
      "\n",
      "**<name of validation method>**\n",
      "- <Provide a detailed description of the validation method, including what it validates, how it performs the validation, and any specific rules or criteria it follows.>\n",
      "\n",
      "**<name of the method that invokes the validation and id generation>**\n",
      "- <Provide a detailed description of the method that invokes both the validation and ID generation methods. Describe its purpose, how it integrates these methods, and any relevant details about its operation.>\n",
      "\n",
      "<End of repeat section for classes.>\n",
      "[/INST]\n",
      "\n",
      "[INST]Additionally, you must include a detailed description of the validation, ID generation, and the method which invokes both of these methods. Ensure clarity and comprehensiveness in these descriptions.[/INST]\n",
      "\n",
      "[INST]If any of the above mentioned field are not applicable, please skip them and not add any unneccessary information[/INST]\n",
      "\n",
      "[INST]Do not add any unnecessary text or information outside of the specified format.[/INST]\n",
      "</s>\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "doc_prompt_template=\"\"\"Prepare a documentation for this java program\n",
    "{javacode} .\"\"\"+doc_format\n",
    "\n",
    "prompt_documentation = PromptTemplate(\n",
    "    input_variables=['javacode'],\n",
    "    template=doc_prompt_template\n",
    ")\n",
    "\n",
    "formatted_prompt = prompt_documentation.format(\n",
    "    javacode=prettify_java_code\n",
    ")\n",
    "\n",
    "print(formatted_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = LLMChain(llm=LLM, prompt=prompt_documentation)\n",
    "raw_docs=chain.run({'javacode':prettify_java_code})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document saved as Java_Program_Documentation.docx\n"
     ]
    }
   ],
   "source": [
    "from docx import Document\n",
    "from docx.shared import RGBColor\n",
    "from docx.enum.text import WD_PARAGRAPH_ALIGNMENT\n",
    "\n",
    "def get_page_width(doc):\n",
    "    # Get the page width based on left and right margins\n",
    "    return doc.sections[0].page_width - doc.sections[0].left_margin - doc.sections[0].right_margin\n",
    "\n",
    "def separate_formatting_rules(doc_string):\n",
    "    parts = doc_string.split('**Formatting Rules**')\n",
    "    main_text = parts[0].strip()\n",
    "    formatting_rules = parts[1].strip() if len(parts) > 1 else \"\"\n",
    "    return main_text, formatting_rules\n",
    "\n",
    "def create_java_program_doc(doc_string, class_diagram_path, file_name='Java_Program_Documentation.docx'):\n",
    "    # Separate the formatting rules from the main text\n",
    "    main_text, formatting_rules = separate_formatting_rules(doc_string)\n",
    "\n",
    "    # Create a new Document\n",
    "    doc = Document()\n",
    "\n",
    "    # Split the main text into lines\n",
    "    lines = main_text.split('\\n')\n",
    "    \n",
    "    for line in lines:\n",
    "        if line.startswith('**Problem Statement**'):\n",
    "            # Problem statement in bold\n",
    "            paragraph = doc.add_paragraph()\n",
    "            run = paragraph.add_run(line.strip(\"**\"))\n",
    "            run.bold = True\n",
    "            run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "        elif line.startswith('**Class Diagram**'):\n",
    "            # Class diagram section\n",
    "            paragraph = doc.add_paragraph()\n",
    "            run = paragraph.add_run(line.strip(\"**\"))\n",
    "            run.bold = True\n",
    "            run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "            # Add class diagram image\n",
    "            page_width = get_page_width(doc)\n",
    "            doc.add_picture(class_diagram_path, width=page_width)\n",
    "        elif line.startswith('**Model Class**'):\n",
    "            # Model class section\n",
    "            paragraph = doc.add_paragraph()\n",
    "            run = paragraph.add_run(line.strip(\"**\"))\n",
    "            run.bold = True\n",
    "            run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "        elif line.startswith('**') and line.endswith('**'):\n",
    "            # Bold class names\n",
    "            paragraph = doc.add_paragraph()\n",
    "            run = paragraph.add_run(line.strip(\"**\"))\n",
    "            run.bold = True\n",
    "            run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "        elif line.startswith('* '):\n",
    "            # Bulleted list item\n",
    "            paragraph = doc.add_paragraph(line.strip(\"* \"), style='List Bullet')\n",
    "            for run in paragraph.runs:\n",
    "                run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "        elif line.startswith('    * '):\n",
    "            # Nested bullet list item\n",
    "            paragraph = doc.add_paragraph(line.strip().replace(\"*\",\"\"), style='List Bullet 2')\n",
    "            for run in paragraph.runs:\n",
    "                run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "        else:\n",
    "            # Regular paragraph\n",
    "            paragraph = doc.add_paragraph(line)\n",
    "            for run in paragraph.runs:\n",
    "                run.font.color.rgb = RGBColor(0, 0, 0)  # Set font color to black\n",
    "\n",
    "    # Save the document\n",
    "    doc.save(file_name)\n",
    "    print(f\"Document saved as {file_name}\")\n",
    "\n",
    "# Example usage\n",
    "doc_string = raw_docs\n",
    "class_diagram_path = \"class_diagram.png\"  # Path to your class diagram PNG file\n",
    "create_java_program_doc(doc_string, class_diagram_path, file_name='Java_Program_Documentation.docx')\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "prompts@infosys",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}