Add web UI

analyze.py

@@ -17,7 +17,8 @@
 ai_client = OpenAIClient(
     base_url="https://api.openai.com/v1/",
     api_key=env.open_ai_key,
-    model="text-embedding-3-small",
+    embeddings_model="text-embedding-3-small",
+    chat_model="gpt-4o"
 )
 
 analyzer = EmbeddingsAnalyzer(embeddings_gateway, chunks_gateway, ai_client)

starter/ai/open_ai_client.py

@@ -1,13 +1,21 @@
+from dataclasses import dataclass
 from typing import List
 
 import requests
 
 
+@dataclass
+class ChatMessage:
+    role: str
+    content: str
+
+
 class OpenAIClient:
-    def __init__(self, base_url: str, api_key: str, model: str):
+    def __init__(self, base_url: str, api_key: str, embeddings_model: str, chat_model: str):
         self.base_url = base_url
         self.api_key = api_key
-        self.model = model
+        self.embeddings_model = embeddings_model
+        self.chat_model = chat_model
 
     def fetch_embedding(self, text) -> List[float]:
         result = requests.post(
@@ -17,10 +25,27 @@ def fetch_embedding(self, text) -> List[float]:
                 "Content-Type": "application/json",
             },
             json={
-                "model": self.model,
+                "model": self.embeddings_model,
                 "input": text,
                 "encoding_format": "float",
             },
         )
 
         return result.json()["data"][0]["embedding"]
+
+    def fetch_chat_completion(self, messages: List[ChatMessage]) -> str:
+        result = requests.post(
+            f"{self.base_url}/chat/completions",
+            headers={
+                "Authorization": f"Bearer {self.api_key}",
+                "Content-Type": "application/json",
+            },
+            json={
+                "model": self.chat_model,
+                "messages": [
+                    {"role": message.role, "content": message.content}
+                    for message in messages
+                ]},
+        )
+
+        return result.json()["choices"][0]["message"]["content"]

starter/app.py

@@ -3,9 +3,16 @@
 import sqlalchemy
 from flask import Flask
 
+from starter.ai.open_ai_client import OpenAIClient
 from starter.database_support.database_template import DatabaseTemplate
+from starter.documents.chunks_gateway import ChunksGateway
+from starter.documents.documents_gateway import DocumentsGateway
 from starter.environment import Environment
 from starter.health_api import health_api
+from starter.index_page import index_page
+from starter.query.query_service import QueryService
+from starter.search.chunks_search_service import ChunksSearchService
+from starter.search.embeddings_gateway import EmbeddingsGateway
 
 logger = logging.getLogger(__name__)
 
@@ -17,6 +24,19 @@ def create_app(env: Environment = Environment.from_env()) -> Flask:
     db = sqlalchemy.create_engine(env.database_url, pool_size=4)
     db_template = DatabaseTemplate(db)
 
+    documents_gateway = DocumentsGateway(db_template)
+    chunks_gateway = ChunksGateway(db_template)
+    embeddings_gateway = EmbeddingsGateway(db_template)
+    ai_client = OpenAIClient(
+        base_url="https://api.openai.com/v1/",
+        api_key=env.open_ai_key,
+        embeddings_model="text-embedding-3-small",
+        chat_model="gpt-4o"
+    )
+    chunks_search_service = ChunksSearchService(embeddings_gateway, chunks_gateway, documents_gateway, ai_client)
+
+    query_service = QueryService(chunks_search_service, ai_client)
+    app.register_blueprint(index_page(query_service))
     app.register_blueprint(health_api(db_template))
 
     return app

starter/database_support/database_template.py

@@ -1,4 +1,4 @@
-from typing import Optional, Any, Callable, TypeVar, Generic
+from typing import Optional, Any, Callable, TypeVar
 
 import sqlalchemy
 from sqlalchemy import Engine, Connection, CursorResult

starter/documents/documents_gateway.py

@@ -37,3 +37,16 @@ def exists(self, source: str, connection: Optional[Connection] = None) -> bool:
         )
 
         return map_one_result(result, lambda row: row["count"] > 0)
+
+    def find(self, id: UUID, connection: Optional[Connection] = None) -> DocumentRecord:
+        result = self.template.query(
+            "select id, source, content from documents where id = :id",
+            connection,
+            id=id,
+        )
+
+        return map_one_result(result, lambda row: DocumentRecord(
+            id=row["id"],
+            source=row["source"],
+            content=row["content"],
+        ))

starter/index_page.py

@@ -0,0 +1,26 @@
+from flask import Blueprint, render_template, request
+from flask.typing import ResponseReturnValue
+
+from starter.query.query_service import QueryService
+
+
+def index_page(query_service: QueryService) -> Blueprint:
+    page = Blueprint('index_page', __name__)
+
+    @page.get('/')
+    def index() -> ResponseReturnValue:
+        return render_template('index.html')
+
+    @page.post('/')
+    def query() -> ResponseReturnValue:
+        user_query = request.form.get('query')
+        result = query_service.fetch_response(user_query)
+
+        return render_template(
+            'response.html',
+            query=user_query,
+            source=result.source,
+            response=result.response,
+        )
+
+    return page

starter/query/query_service.py

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+
+from starter.ai.open_ai_client import OpenAIClient, ChatMessage
+from starter.search.chunks_search_service import ChunksSearchService
+
+
+@dataclass
+class QueryResult:
+    source: str
+    response: str
+
+
+class QueryService:
+    def __init__(self, chunks_search_service: ChunksSearchService, ai_client: OpenAIClient):
+        self.chunks_search_service = chunks_search_service
+        self.ai_client = ai_client
+
+    def fetch_response(self, query: str) -> QueryResult:
+        chunk = self.chunks_search_service.search_for_relevant_chunk(query)
+        response = self.ai_client.fetch_chat_completion([
+            ChatMessage(role="system", content="You are a reporter for a major world newspaper."),
+            ChatMessage(role="system", content="Write your response as if you were writing a short, high-quality news"
+                                               "article for your paper. Limit your response to one paragraph."),
+            ChatMessage(role="system", content=f"Use the following article for context: {chunk.content}"),
+            ChatMessage(role="user", content=query),
+        ])
+
+        return QueryResult(
+            source=chunk.source,
+            response=response
+        )

starter/search/chunks_search_service.py

@@ -0,0 +1,35 @@
+from dataclasses import dataclass
+
+from starter.ai.open_ai_client import OpenAIClient
+from starter.documents.chunks_gateway import ChunksGateway
+from starter.documents.documents_gateway import DocumentsGateway
+from starter.search.embeddings_gateway import EmbeddingsGateway
+
+
+@dataclass
+class ChunkSearchResult:
+    content: str
+    source: str
+
+
+class ChunksSearchService:
+    def __init__(self,
+                 embeddings_gateway: EmbeddingsGateway,
+                 chunks_gateway: ChunksGateway,
+                 documents_gateway: DocumentsGateway,
+                 open_ai_client: OpenAIClient):
+        self.embeddings_gateway = embeddings_gateway
+        self.chunks_gateway = chunks_gateway
+        self.documents_gateway = documents_gateway
+        self.open_ai_client = open_ai_client
+
+    def search_for_relevant_chunk(self, query: str) -> ChunkSearchResult:
+        vector = self.open_ai_client.fetch_embedding(query)
+        chunk_id = self.embeddings_gateway.find_similar_chunk_id(vector)
+        chunk = self.chunks_gateway.find(chunk_id)
+        document = self.documents_gateway.find(chunk.document_id)
+
+        return ChunkSearchResult(
+            content=chunk.content,
+            source=document.source,
+        )

starter/search/embeddings_gateway.py

@@ -5,6 +5,7 @@
 
 from starter.database_support.database_template import DatabaseTemplate
 from starter.database_support.result_mapping import map_results, map_one_result
+from starter.search.vector_support import vector_to_string
 
 
 class EmbeddingsGateway:
@@ -33,7 +34,7 @@ def find_similar_chunk_id(self, vector: List[float], connection: Optional[Connec
         result = self.template.query(
             """select e.chunk_id from embeddings e order by e.embedding <=> :vector limit 1""",
             connection,
-            vector=vector,
+            vector=vector_to_string(vector),
         )
 
         return map_one_result(result, lambda row: row["chunk_id"])

starter/search/vector_support.py

@@ -0,0 +1,5 @@
+from typing import List
+
+
+def vector_to_string(vector: List[float]) -> str:
+    return "[" + ",".join([str(v) for v in vector]) + "]"