Revamp the User Interface

Implement AI streaming
Implement AI features for single author
Add the ability to download single author information.

ai/.env

@@ -1,3 +1,3 @@
 REDIS_BROKER=redis://localhost:6379/0
-REDIS_RESULT=redis://localhost:6379/0
-MISTRAL_API_KEY=hSZCZCDRWdhsfd9N2XZWEIWNeeMB5trf
+REDIS_BACKEND=redis://localhost:6379/0
+MISTRAL_API_KEY=cCpcQhIsggTkErjaa3ZVCtu6XAjetH0h

ai/Dockerfile

@@ -0,0 +1,28 @@
+FROM python:3.12-slim-bullseye AS builder
+
+RUN pip install poetry==1.8.3
+
+ENV POETRY_NO_INTERACTION=1 \
+    POETRY_VIRTUALENVS_IN_PROJECT=1 \
+    POETRY_VIRTUALENVS_CREATE=1 \
+    POETRY_CACHE_DIR=/tmp/poetry_cache
+
+WORKDIR /worker
+
+COPY pyproject.toml poetry.lock ./
+RUN touch README.md
+
+RUN --mount=type=cache,target=$POETRY_CACHE_DIR poetry install --no-root
+
+FROM python:3.12-slim-bullseye AS runtime
+
+ENV VIRTUAL_ENV=/worker/.venv \
+    PATH="/worker/.venv/bin:$PATH"
+
+COPY --from=builder ${VIRTUAL_ENV} ${VIRTUAL_ENV}
+
+COPY . ./worker
+
+WORKDIR /worker
+
+ENTRYPOINT ["celery", "-A", "celery_app.celery", "worker", "--loglevel=info", "--pool=solo", "-Q", "ai"]

ai/app.py

@@ -1,14 +1,35 @@
 import asyncio
 from io import StringIO
 import json
+import time
 
 from celery import current_app as celery
+from celery import Task
 from mistralapi import main as task_ai
+from mistralapi import async_main as async_task_ai
+from redis import Redis
+
+redis = Redis(host='redis', port=6379, db=1)
 
 @celery.task
 def process_model(data: dict):
     input_data = json.dumps(data)
     with StringIO(input_data) as in_buffer, StringIO() as out_buffer:
         task_ai(in_buffer, out_buffer)
         output_data = out_buffer.getvalue()
+    return {'text': output_data}
+
+@celery.task(bind=True)
+def process_model_two(self: Task, data:dict):
+    input_data = json.dumps(data)
+    v = 0
+    with StringIO(input_data) as in_buffer, StringIO() as out_buffer:
+        for chunk in async_task_ai(in_buffer, out_buffer):
+            print(f"Chunk {chunk}")
+            resp = redis.xadd(f"ai_completions_{self.request.id}", {b"data": chunk.encode(), "ts": time.time(), 'v': v})
+            v += 1
+            print(resp, v)
+        output_data = out_buffer.getvalue()
+        resp = redis.xadd(f"ai_completions_{self.request.id}", {b"data": "[END]".encode(), "ts": time.time(), 'v': v})
+        print(resp, v)
     return {'text': output_data}

ai/async_test.py

@@ -0,0 +1,37 @@
+import os
+from mistralai import Mistral
+import asyncio
+import text_io
+from io import StringIO
+
+from dotenv import load_dotenv
+load_dotenv()
+
+# Ensure the os module is imported
+api_key = os.environ["MISTRAL_API_KEY"]
+model = "open-mistral-nemo-2407"
+
+client = Mistral(api_key=api_key)
+
+
+async def main():
+    string = StringIO()
+    # Calling the async method inside an async function
+    async_response = client.chat.stream(
+        model=model,
+        messages=[
+            {
+                "role": "user", 
+                "content": "Who is the best French painter?",
+            },
+        ]
+    )
+
+    # Using async for loop inside an async function
+
+    for chunk in async_response: 
+        print(chunk.data.choices[0].delta.content, end="", flush=True)
+        #text_io.write_paragraph_to_file(chunk.data.choices[0].delta.content)
+
+# Run the async function using asyncio
+asyncio.run(main())

ai/celery_app.py

@@ -1,8 +1,11 @@
 from celery import Celery
 from app import *
 
+from dotenv import load_dotenv
+load_dotenv()
+import os
 
-celery = Celery("ai", broker="redis://localhost:6379/0", backend="redis://localhost:6379/0")
+celery = Celery("ai", broker=os.environ["REDIS_BROKER"], backend=os.environ["REDIS_BACKEND"])
 celery.conf.broker_connection_retry_on_startup = True
 celery.conf.task_routes = {
     'app.process_model': {'queue': 'ai'},

ai/mistralapi.py

@@ -20,6 +20,9 @@
 }
 
 def get_api_key():
+    if not os.environ.get("MISTRAL_API_KEY", None):
+        with open("/run/secrets/mistral_ai_key", 'r') as f:
+            return f.read().strip()
     return os.environ.get("MISTRAL_API_KEY")
 
 def initialize_client(api_key):
@@ -30,7 +33,9 @@ def construct_prompt(publications_file_path):
         "Summarize the following publications of the authors. Analyze their expertise, "
         "the significance of the studies, and provide insights based on the titles and venues. "
         "The summary should be detailed, comprehensive, and suitable for presentation. "
-        "Use paragraphs and avoid repetition.\n\n"
+        "Use paragraphs and avoid repetition. "
+        "Use headings to separate the summaries of different authors. "
+        "\n\n"
     )
 
     publication_data = text_io.read_large_file(publications_file_path)
@@ -62,5 +67,28 @@ def main(in_buffer: StringIO, out_buffer: StringIO):
     # save_summary_to_file(summary, out_buffer)
     out_buffer.write(summary)
 
+def async_main(in_buffer: StringIO, out_buffer: StringIO):
+    selected_model_name = "Mistral Nemo"  # Change this to select a different model
+    model_name = MODEL_OPTIONS[selected_model_name]
+
+    api_key = get_api_key()
+    if not api_key:
+        raise ValueError("MISTRAL_API_KEY not set in .env file.")
+
+    client = Mistral(api_key=api_key)
+    prompt = construct_prompt(in_buffer)
+    completion = client.chat.stream(
+        max_tokens=2**32,
+        model=model_name,
+        messages=[
+            {"role": "user", "content": prompt}
+        ],
+    )
+    for chunk in completion:
+        content = chunk.data.choices[0].delta.content
+        out_buffer.write(content)
+        yield content
+
+
 if __name__ == "__main__":
     main()

ai/pyproject.toml

@@ -10,6 +10,8 @@ readme = "README.md"
 python = "^3.12"
 celery = "^5.4.0"
 mistralai = "^1.0.2"
+python-dotenv = "^1.0.1"
+redis = "^5.0.8"
 
 
 [build-system]

client/src/components/AIModel.tsx

@@ -1,20 +1,18 @@
 import { API_SERVER } from "../config/config";
-import usePolling from "../hooks/usePolling";
 import Markdown from 'react-markdown'
+import useStream from "../hooks/useStream";
 
 export default function AIModel({aiTaskId}: {aiTaskId: string | null}) {
   const {
     data: aiData,
     loading,
     error,
-  } = usePolling<{status: "PENDING" | "FAILURE" | "SUCCESS"} | {text: string}>(`${API_SERVER}/tasks/ai/${aiTaskId}/result`, 1500, {
-    enabled: aiTaskId?.length !== 0,
-  });
+  } = useStream<string>(`${API_SERVER}/tasks/ai/${aiTaskId}/result`);
   console.log(aiData)
 
     return (
     <div className="flex flex-col h-full w-full">
-        {aiData && ("text" in aiData) && <Markdown>{aiData.text}</Markdown>}
+        {aiData && <Markdown>{aiData}</Markdown>}
         {(!aiData || loading) && <p>Loading...</p>}
         {error && <p>Error: {error.message}</p>}
     </div>)

client/src/components/Button.tsx

@@ -1,6 +1,6 @@
 export default function Button({children, onClick, className}: {children: React.ReactNode; onClick: () => void; className?: string}) {
     return (
-        <button onClick={onClick} className={`transition-all flex items-center gap-2 flex-row text-black dark:text-white rounded-xl outline dark:outline-neutral-700 dark:hover:outline-neutral-600 dark:active:bg-neutral-600 px-4 py-2 ${className}`}>
+        <button onClick={onClick} className={`transition-all flex items-center gap-2 flex-row text-black dark:text-white rounded-xl outline outline-neutral-600 hover:outline-neutral-800 active:bg-neutral-800 active:text-white dark:outline-neutral-700 dark:hover:outline-neutral-600 dark:active:bg-neutral-600 px-4 py-2 ${className}`}>
             {children}
         </button>
     )

client/src/components/PreviewTable.tsx

@@ -1,8 +1,8 @@
 export default function PreviewTable({headers, data, footer}:{headers: string[], data: string[][], footer: string}) {
     return (
         <div className="flex bg-neutral-200 dark:bg-neutral-900 rounded-3xl outline outline-2 outline-neutral-800 text-black dark:outline-neutral-400 gap-4 dark:text-white">
-            <div className="overflow-x-auto">
-                <table className="table-auto divide-y-2 divide-neutral-800 dark:divide-neutral-400">
+            <div className="overflow-x-auto w-full table-fixed">
+                <table className="table-auto w-full divide-y-2 divide-neutral-800 dark:divide-neutral-400">
                     <thead className="">
                         <tr className="divide-x-2 divide-neutral-700 dark:divide-neutral-400">
                             {headers.map((header, index) => <th key={index} className="px-4 py-2">{header}</th>)}
@@ -13,7 +13,7 @@ export default function PreviewTable({headers, data, footer}:{headers: string[],
                             row.length === 6 && <tr key={index} className="dark:odd:bg-neutral-950 odd:bg-neutral-50 divide-x-2 dark:divide-neutral-800 divide-neutral-700 dark:even:bg-neutral-900 even:bg-neutral-100">
                                 <td className="px-4 py-2 w-fit max-w-sm">{row[0]}</td>
                                 <td className="px-4 py-2 w-fit max-w-sm"><a href={row[4]} className="text-blue-500 hover:underline">{row[1]}</a></td>
-                                <td className="px-4 py-2 w-fit max-w-sm">{row[2]}<br />({row[5]})</td>
+                                <td className="px-4 py-2 w-fit max-w-sm">{row[2]}<br /><span className="text-neutral-400">({row[5]})</span></td>
                                 <td className="px-4 py-2 w-fit max-w-sm">{row[3]}</td>
                             </tr>
                         ))}

client/src/components/Search.tsx

@@ -28,16 +28,17 @@ export default function Search({
         onChange={(e) => setQuery(e.target.value)}
         className="h-full flex-1 bg-transparent text-black dark:text-white outline-none dark:placeholder:text-neutral-400 placeholder:text-neutral-600"
         type="search"
+        aria-description="Search for a person..."
         placeholder="Search for a person..."
       />
       <div className="flex gap-2">
-        <button onClick={filterIconPressed}>
+        <button aria-description="Filter records" onClick={filterIconPressed}>
           <FilterIcon
             filled={filterEnabled}
             className="h-full size-6 text-neutral-600 hover:text-black dark:text-neutral-400 dark:hover:text-white"
           />
         </button>
-        <button onClick={uploadIconPressed}>
+        <button aria-description="Upload excel file" onClick={uploadIconPressed}>
           <ArrowUpTrayIcon className="h-full size-6 text-neutral-600 hover:text-black dark:text-neutral-400 dark:hover:text-white" />
         </button>
         <button className="hidden" onClick={goIconPressed}>

client/src/hooks/usePolling.ts

@@ -4,7 +4,7 @@ export default function usePolling<T>(
   url: string,
   delay_msec: number,
   pollingOpts: { enabled?: boolean } = { enabled: true }
-): { data: T | null; loading: boolean; error: Error | null, resetData: () => void } {
+): { data: T | null; loading: boolean; error: Error | null, resetData: () => void, headers: Headers | null } {
   // final constants
   const [data, setData] = useState<T | null>(null);
   const [loading, setLoading] = useState<boolean>(false);
@@ -13,6 +13,7 @@ export default function usePolling<T>(
   const [isErrored, setIsErrored] = useState<boolean>(false);
   const [dataLoaded, setDataLoaded] = useState<boolean>(false);
   const timerRef = useRef<NodeJS.Timeout | null>(null);
+  const [headers, setHeaders] = useState<Headers | null>(null);
 
   const fetchData = useCallback(async () => {
     if (!pollingOpts.enabled) {
@@ -27,6 +28,7 @@ export default function usePolling<T>(
         setData(json);
         setLoading(false);
         clearInterval(timerRef.current ?? undefined);
+        setHeaders(res.headers);
         return;
       }
       if (json["status"] === "SUCCESS") {
@@ -64,5 +66,5 @@ export default function usePolling<T>(
     setData(null);
   }
 
-  return { loading, data, error, resetData };
+  return { loading, data, error, resetData, headers };
 }

client/src/hooks/useStream.ts

@@ -0,0 +1,51 @@
+import { useEffect, useState, useRef } from 'react';
+export default function useStream<T = string>(url:string) {
+    const [data, setData] = useState<T | null>(null);
+    const [loading, setLoading] = useState(true);
+    const [errorObject, setErrorObject] = useState<Error | null>(null);
+    const text = useRef("");
+
+    useEffect(() => {
+        const controller = new AbortController();
+        const run = async (controller:AbortController) => {   
+          setLoading(true)
+          let response
+          try {
+
+            response = await fetch(url, {
+                method: 'GET',
+                signal: controller.signal,
+            })
+
+          } catch (error) {
+            if ((error as Error).name === "AbortError") {
+              return;
+            }
+            setErrorObject(error as Error)
+          }
+          setLoading(false);
+
+          const reader = response!.body!.getReader();
+          const chunks = [];
+          let done, value;
+          const dec = new TextDecoder()
+
+          while (!done) {
+            ({ value, done } = await reader.read());
+            if (done) {
+              return chunks;
+            }
+            const strval = dec.decode(value, { stream: true })
+            console.log(chunks, strval, text.current)
+            chunks.push(strval);
+            text.current += strval
+            setData(text.current as T)      
+          }
+        }
+        run(controller)
+        return () => controller.abort();
+
+    }, [url]);
+
+    return { data, loading, error: errorObject};
+}