refactor(service): update LexAIService to use new formatter utilities

- Improved module and class docstrings for clarity and professionalism
- Integrated `format_legal_response` and `format_references` to standardize HTML formatting
- Enhanced maintainability by separating formatting logic from business logic

lexai/services/lexai_service.py

@@ -1,23 +1,56 @@
 """
-Service layer for handling LexAI application logic.
+LexAI service layer for handling user queries.
 
-Provides an abstraction between the UI and the core matching engine.
+This module defines a service class that interfaces with the core match engine,
+processes the results, and formats them for display in the UI.
 """
 
 from lexai.core.match_engine import generate_matches
+from lexai.ui.formatters import format_legal_response, format_references
 
 
 class LexAIService:
+    """
+    Service class that handles legal queries by invoking the match engine and
+    formatting the results for UI presentation.
+    """
+
     @staticmethod
     def handle_query(query: str, location: str) -> str:
         """
-        Process a user query by invoking the match engine.
+        Handles a user query and returns an HTML-formatted response.
+
+        This method:
+        - Calls the semantic match engine with the given query and location.
+        - Extracts both the AI-generated response and the list of matched legal entries.
+        - Constructs an HTML string that includes the AI's response followed by
+          a reference list linking to legal documents.
 
-        Parameters:
-            query (str): The legal question from the user.
-            location (str): The jurisdiction for the query.
+        Parameters
+        ----------
+        query : str
+            The legal question asked by the user.
+        location : str
+            The jurisdiction to search within.
 
-        Returns:
-            str: HTML-formatted response with AI-generated content and references.
+        Returns
+        -------
+        str
+            A formatted HTML string with the AI response and relevant matches.
         """
-        return generate_matches(query, location)
+        result = generate_matches(query, location)
+
+        gpt_response = result.get("response", "").strip()
+        matches = result.get("matches", [])
+
+        if (
+            not isinstance(matches, list)
+            or not matches
+            or not isinstance(matches[0], dict)
+        ):
+            return format_legal_response(gpt_response or "No matches found.")
+
+        return (
+            format_legal_response(gpt_response) +
+            format_references(matches)
+        )

lexai/ui/formatters.py

@@ -11,6 +11,16 @@ from html import escape
 def format_legal_response(response_text: str) -> str:
     """
     Wrap the AI-generated legal response in HTML for UI rendering.
+
+    Parameters
+    ----------
+    response_text : str
+        The main response text from the assistant.
+
+    Returns
+    -------
+    str
+        HTML-formatted string with a 'Response' header and the content.
     """
     return (
         "<p><strong>Response:</strong></p>"
@@ -21,18 +31,29 @@ def format_legal_response(response_text: str) -> str:
 def format_references(matches: list[dict]) -> str:
     """
     Format a list of top document matches into an HTML reference list.
+
+    Parameters
+    ----------
+    matches : list of dict
+        List of matched legal documents, each containing 'url', 'title', and 'subtitle'.
+
+    Returns
+    -------
+    str
+        HTML-formatted reference section with clickable links.
     """
     if not matches:
         return "<p><strong>References:</strong> None found.</p>"
 
-    reference_html = "<p><strong>References:</strong></p><ul>"
+    html = "<p><strong>References:</strong></p><ul>"
     for match in matches:
         url = escape(match.get("url", "#"))
         title = escape(match.get("title", "Untitled"))
         subtitle = escape(match.get("subtitle", ""))
-        reference_html += (
-            f'<li><a href="{url}" target="_blank" rel="noopener noreferrer">'
+        html += (
+            "<li>"
+            f"<a href=\"{url}\" target=\"_blank\" rel=\"noopener noreferrer\">"
             f"{title}: {subtitle}</a></li>"
         )
-    reference_html += "</ul>"
-    return reference_html
+    html += "</ul>"
+    return html

lexai/ui/gradio_interface.py

@@ -15,10 +15,25 @@ from lexai.services.lexai_service import LexAIService
 logger = logging.getLogger(__name__)
 
 APP_DESCRIPTION = """
-LexAI is an AI-powered legal assistant that uses GPT-4 and semantic search
-to provide jurisdiction-specific legal guidance.
+LexAI is an AI-powered legal assistant that provides jurisdiction-specific guidance.
+It combines GPT-4 with semantic search to retrieve relevant legal information quickly.
 """
 
+DISCLAIMER_TEXT = """
+<div style='text-align: center; font-size: 0.9em; color: gray; margin-top: 1em;'>
+Results may be inaccurate. Always verify with a legal professional.
+</div>
+"""
+
+EXAMPLE_QUERIES = [
+    ["Is building a rock cairn outdoors allowed by law?", "Boulder"],
+    ["Can I legally possess a dog as a pet?", "Denver"],
+    ["Am I allowed to go shirtless in public?", "Boulder"],
+    ["What is the max legal height for a structure?", "Denver"],
+    ["Is indoor furniture on porches allowed?", "Boulder"],
+    ["Can I graze llamas on public land?", "Denver"],
+]
+
 
 def build_interface():
     """
@@ -55,9 +70,7 @@ def build_interface():
                 gr.Button("Flag", variant="secondary")
 
         def handle_submit(query, location):
-            return gr.update(
-                value=LexAIService.handle_query(query, location)
-            )
+            return gr.update(value=LexAIService.handle_query(query, location))
 
         def handle_clear():
             return gr.update(value="Response will appear here.")
@@ -73,16 +86,11 @@ def build_interface():
         )
 
         gr.Examples(
-            examples=[
-                ["Is building a rock cairn outdoors allowed by law?", "Boulder"],
-                ["Can I legally possess a dog as a pet?", "Denver"],
-                ["Am I allowed to go shirtless in public?", "Boulder"],
-                ["What is the max legal height for a structure?", "Denver"],
-                ["Is indoor furniture on porches allowed?", "Boulder"],
-                ["Can I graze llamas on public land?", "Denver"],
-            ],
+            examples=EXAMPLE_QUERIES,
             inputs=[query_input, location_input]
         )
 
+        gr.HTML(DISCLAIMER_TEXT)
+
     logger.info("LexAI interface built.")
     return iface

tests/test_matcher.py

@@ -14,14 +14,13 @@ from lexai.core.matcher import find_top_matches
 
 @pytest.fixture
 def sample_embeddings():
-    """Sample array of 5 embedding vectors."""
+    """Sample array of 4 embedding vectors."""
     return np.array(
         [
             [1.0, 0.1, 0.1],
             [0.8, 0.3, 0.2],
             [0.5, 0.5, 0.5],
             [0.1, 0.1, 1.0],
-            [0.0, 0.0, 0.0],
         ],
         dtype=np.float32,
     )

tests/test_openai_client.py

@@ -1,39 +1,99 @@
 """
-Tests for the OpenAI client service in lexai.services.openai_client.
+OpenAI client functions for embedding generation and GPT-4 completions.
+
+This module provides utilities to interact with OpenAI’s API, including
+embedding generation and chat-based completion using the configured models.
 """
 
-from unittest.mock import MagicMock, patch
+import os
 
 import numpy as np
+from openai import OpenAI
+from openai.types.chat import ChatCompletion
+from openai.types.embedding import Embedding
+
+from lexai.config import (
+    EMBEDDING_MODEL,
+    GPT4_FREQUENCY_PENALTY,
+    GPT4_MAX_TOKENS,
+    GPT4_MODEL,
+    GPT4_PRESENCE_PENALTY,
+    GPT4_TEMPERATURE,
+    GPT4_TOP_P,
+)
+
+
+def get_client() -> OpenAI:
+    """
+    Returns a new instance of the OpenAI client using the current API key.
+
+    Returns
+    -------
+    OpenAI
+        An authenticated OpenAI client.
+    """
+    api_key = os.getenv("OPENAI_API_KEY")
+    if not api_key:
+        raise EnvironmentError(
+            "OPENAI_API_KEY environment variable is not set.")
+    return OpenAI(api_key=api_key)
+
 
-from lexai.services.openai_client import get_chat_completion, get_embedding
+def get_embedding(text: str) -> np.ndarray:
+    """
+    Generates a numerical embedding for the provided text using OpenAI's model.
 
+    Parameters
+    ----------
+    text : str
+        The input text to embed.
 
-@patch("lexai.services.openai_client.client")
-def test_get_embedding_success(mock_client):
-    """Test that get_embedding returns the correct NumPy array."""
-    mock_response = MagicMock()
-    mock_response.data = [MagicMock(embedding=[0.1, 0.2, 0.3])]
-    mock_client.embeddings.create.return_value = mock_response
+    Returns
+    -------
+    np.ndarray
+        The embedding vector as a NumPy array.
+    """
+    client = get_client()
+    response: Embedding = client.embeddings.create(
+        input=text, model=EMBEDDING_MODEL)
+    return np.array(response.data[0].embedding)
 
-    embedding = get_embedding("Test input")
-    assert isinstance(embedding, np.ndarray)
-    np.testing.assert_array_equal(embedding, np.array([0.1, 0.2, 0.3]))
 
+def get_chat_completion(
+    role_description: str,
+    jurisdiction_summary: str,
+    query: str,
+) -> str:
+    """
+    Generates a GPT-4 response based on the user’s query and legal context.
 
-@patch("lexai.services.openai_client.client")
-def test_get_chat_completion_success(mock_client):
-    """Test that get_chat_completion returns the expected string."""
-    mock_choice = MagicMock()
-    mock_choice.message.content = "Here is your legal summary."
-    mock_response = MagicMock()
-    mock_response.choices = [mock_choice]
-    mock_client.chat.completions.create.return_value = mock_response
+    Parameters
+    ----------
+    role_description : str
+        Describes the assistant's role and intended tone or expertise.
+    jurisdiction_summary : str
+        A stringified summary of relevant legal documents or search results.
+    query : str
+        The user's legal question.
 
-    response = get_chat_completion(
-        role_description="You are a legal assistant.",
-        context_summary="1. Case A\n2. Case B",
-        query="What is the precedent for X?"
+    Returns
+    -------
+    str
+        The assistant's response.
+    """
+    client = get_client()
+    response: ChatCompletion = client.chat.completions.create(
+        model=GPT4_MODEL,
+        messages=[
+            {"role": "system", "content": role_description.strip()},
+            {"role": "system", "content": jurisdiction_summary.strip()},
+            {"role": "user", "content": query.strip()},
+            {"role": "assistant", "content": ""},
+        ],
+        temperature=GPT4_TEMPERATURE,
+        max_tokens=GPT4_MAX_TOKENS,
+        top_p=GPT4_TOP_P,
+        frequency_penalty=GPT4_FREQUENCY_PENALTY,
+        presence_penalty=GPT4_PRESENCE_PENALTY,
     )
-    assert isinstance(response, str)
-    assert response == "Here is your legal summary."
+    return response.choices[0].message.content.strip()