Update app.py

app.py

@@ -3,21 +3,22 @@ import json
 import yaml
 import httpx
 from transformers import pipeline
-import fitz  # PyMuPDF
-import docx  # python-docx for handling doc/docx files
+import fitz  # PyMuPDF for PDF
+import docx  # python-docx for DOCX
 import logging
 
-# Initialize the Hugging Face model pipeline (using text generation model for simplicity)
+# Initialize the Hugging Face model pipeline (using a text generation model for simplicity)
 llm_model = pipeline('text-generation', model="bigscience/bloom-560m")
 
 # Configure logging
 logging.basicConfig(filename='api_client.log', level=logging.INFO, 
                     format='%(asctime)s - %(levelname)s - %(message)s')
 
+
 def parse_api_spec(api_spec_content):
     """
-    This function parses the uploaded API specification and returns endpoints and their parameters.
-    Assumes OpenAPI format.
+    Parses the uploaded API specification and returns endpoints and their parameters.
+    Assumes OpenAPI format (JSON/YAML).
     """
     try:
         # Try to load as JSON first
@@ -71,15 +72,36 @@ class APIClient:
     """
     return interface_code
 
+
 def use_llm_to_extract(api_spec):
     """
-    Uses a Hugging Face model to extract API details.
+    Uses a Hugging Face model to extract all API methods, parameters, and response details from the given API spec.
     """
-    prompt = f"Extract API endpoints and parameters from the following API spec:\n{api_spec}"
-    response = llm_model(prompt, max_length=200, num_return_sequences=1)
+    prompt = f"""
+    Please extract all the API methods, their respective endpoints, parameters (both required and optional), and response details from the following API specification:
+    
+    {api_spec}
+    
+    The extracted details should include:
+    1. HTTP Method (e.g., GET, POST, PUT, DELETE)
+    2. API Endpoint (e.g., /api/v1/recharge-cash-account)
+    3. Input Parameters:
+        - Parameter name
+        - Parameter type (e.g., String, Integer)
+        - Parameter description (if provided)
+        - Whether the parameter is mandatory (M) or optional (O)
+    4. Output (Response) Parameters:
+        - Parameter name
+        - Parameter type (e.g., String, Integer)
+        - Parameter description (if provided)
+    5. Example Request and Example Response
     
+    Provide a clear summary of each API method and its details.
+    """
+    response = llm_model(prompt, max_length=500, num_return_sequences=1)
     return response[0]['generated_text']
 
+
 def read_pdf(file):
     """
     Extracts text from PDF file using PyMuPDF (fitz).
@@ -90,6 +112,7 @@ def read_pdf(file):
         text += page.get_text()
     return text
 
+
 def read_docx(file):
     """
     Extracts text from DOC/DOCX file using python-docx.
@@ -98,6 +121,7 @@ def read_docx(file):
     text = "\n".join([para.text for para in doc.paragraphs])
     return text
 
+
 def main():
     st.title("API Spec Uploader and Python Interface Generator")
 
@@ -106,50 +130,47 @@ def main():
 
     if uploaded_file is not None:
         file_type = uploaded_file.type
-    
-        # Handle different file types
+
         if file_type == "application/json":
             api_spec_content = uploaded_file.read().decode("utf-8")
         elif file_type == "application/x-yaml":
             api_spec_content = uploaded_file.read().decode("utf-8")
         elif file_type == "application/pdf":
+            # Extract text from PDF
             api_spec_content = read_pdf(uploaded_file)
         elif file_type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document" or file_type == "application/octet-stream":
-            # Handle docx or unknown types
+            # Handle DOCX or unknown types
             api_spec_content = read_docx(uploaded_file)
         else:
-            st.error(f"Unsupported file format: {file_type}")
+            st.error("Unsupported file format.")
             return
 
-        # Display API Spec
-        st.subheader("Uploaded API Specification")
-        st.code(api_spec_content, language="json" if file_type == "application/json" else "yaml")
-
         # Extract API endpoints and parameters using LLM
         with st.spinner('Extracting API information using Hugging Face LLM...'):
             extracted_info = use_llm_to_extract(api_spec_content)
             st.subheader("Extracted Information from LLM")
             st.write(extracted_info)
 
-        # Parse the API spec manually to display extracted endpoints
-        endpoints = parse_api_spec(api_spec_content)
-        
-        if endpoints:
-            st.subheader("Parsed Endpoints and Parameters")
-            for endpoint, params in endpoints.items():
-                st.write(f"**{endpoint}**")
-                st.json(params)
-
-            # Generate Python interface code with logging
-            python_interface_code = generate_python_interface(endpoints)
-
-            st.subheader("Generated Python Interface with Logging")
-            st.code(python_interface_code, language="python")
+        # Parse the API spec manually to display extracted endpoints (if JSON or YAML)
+        if file_type in ["application/json", "application/x-yaml"]:
+            endpoints = parse_api_spec(api_spec_content)
             
-            st.download_button(label="Download Python Interface",
-                               data=python_interface_code,
-                               file_name="api_interface.py",
-                               mime="text/plain")
+            if endpoints:
+                st.subheader("Parsed Endpoints and Parameters")
+                for endpoint, params in endpoints.items():
+                    st.write(f"**{endpoint}**")
+                    st.json(params)
+
+                # Generate Python interface code with logging
+                python_interface_code = generate_python_interface(endpoints)
+
+                st.subheader("Generated Python Interface with Logging")
+                st.code(python_interface_code, language="python")
+                
+                st.download_button(label="Download Python Interface",
+                                data=python_interface_code,
+                                file_name="api_interface.py",
+                                mime="text/plain")
 
 if __name__ == '__main__':
     main()