From a5ccbe7b32c2b2197858c8c5663625a1b2e1561b Mon Sep 17 00:00:00 2001
From: taranbeer <staranbeer@gmail.com>
Date: Thu, 12 Oct 2023 15:49:42 +0530
Subject: [PATCH] minor fixes

---
 .../inference/_components/promptInference.ts  |  1 +
 .../inference/_components/search-output.tsx   | 27 ++++++++++---------
 .../_components/useInferenceSession.ts        | 19 +++++++++----
 3 files changed, 29 insertions(+), 18 deletions(-)

diff --git a/app/(services)/inference/_components/promptInference.ts b/app/(services)/inference/_components/promptInference.ts
index 6e62ed4..f3da928 100644
--- a/app/(services)/inference/_components/promptInference.ts
+++ b/app/(services)/inference/_components/promptInference.ts
@@ -25,6 +25,7 @@ export const promptInference = async (
       "text",
       (chunk) => {
         const parsed = JSON.parse(chunk) as { choices: { text: string }[] }
+        console.log(parsed)
         const text = parsed.choices[0].text
         response += text
         handleNewChunk(text)
diff --git a/app/(services)/inference/_components/search-output.tsx b/app/(services)/inference/_components/search-output.tsx
index 747ed1f..d1d3785 100644
--- a/app/(services)/inference/_components/search-output.tsx
+++ b/app/(services)/inference/_components/search-output.tsx
@@ -3,23 +3,19 @@
 import React, { useCallback, useEffect, useState } from "react"
 import { useStore } from "@/store/store"
 import { Text } from "@components/text"
-import {
-  Aperture,
-  RefreshCcw,
-  ThumbsDown,
-  ThumbsUp,
-  UserCircle2Icon,
-} from "lucide-react"
+import { Aperture, RefreshCcw, UserCircle2Icon } from "lucide-react"
 
 import { useToast } from "@/components/ui/use-toast"
 import { CometSession } from "@/components/comet-search"
 import { LoadingDots } from "@/components/comet-search/icons"
 import { MarkdownParser } from "@/components/markdown/markdown-parser"
 
+import { InferenceSession } from "./useInferenceSession"
+
 export const SearchOutput = React.forwardRef<HTMLDivElement, any>(
   (
     props: {
-      session: CometSession
+      session: InferenceSession
       streamMessage: string | undefined
       isLoading: boolean
       error: string | undefined
@@ -53,11 +49,16 @@ export const SearchOutput = React.forwardRef<HTMLDivElement, any>(
                   text={i.text}
                 />
               ) : (
-                <SearchReply
-                  key={`${i.conversationId}/${i.from}`}
-                  text={i.text}
-                  conversationId={i.conversationId}
-                />
+                <div>
+                  <SearchReply
+                    key={`${i.conversationId}/${i.from}`}
+                    text={i.text}
+                    conversationId={i.conversationId}
+                  />
+                  {/* <div className="flex items-center gap-5 bg-subdued px-14 pb-4 text-soft">
+                    {i?.latency} {i?.tokens}
+                  </div> */}
+                </div>
               )
             )}
           {props?.streamMessage && (
diff --git a/app/(services)/inference/_components/useInferenceSession.ts b/app/(services)/inference/_components/useInferenceSession.ts
index 9b3ef3c..d7f1c31 100644
--- a/app/(services)/inference/_components/useInferenceSession.ts
+++ b/app/(services)/inference/_components/useInferenceSession.ts
@@ -1,16 +1,22 @@
 import { Dispatch, SetStateAction, useCallback, useRef, useState } from "react"
+import dayjs from "dayjs"
+import relativeTime from "dayjs/plugin/relativeTime"
 
 import { useToast } from "@/components/ui/use-toast"
 
 import { promptInference } from "./promptInference"
 import { GlobalInference } from "./search"
 
-type Session = {
+dayjs.extend(relativeTime)
+
+export type InferenceSession = {
   sessionId: string
   messages: {
     from: "agent" | "human"
     text: string
     conversationId?: string | null
+    latency?: string
+    tokens?: number
   }[]
 }
 
@@ -21,7 +27,7 @@ export function useInferenceSession(
   globalInference: GlobalInference,
   setGlobalInference: Dispatch<SetStateAction<GlobalInference>>
 ) {
-  const [session, setSession] = useState({} as Session)
+  const [session, setSession] = useState({} as InferenceSession)
   const [isLoading, setIsLoading] = useState(false)
   const [isDisabled, setIsDisabled] = useState(false)
   const { toast } = useToast()
@@ -29,7 +35,7 @@ export function useInferenceSession(
 
   const [error, setError] = useState<any>()
   const resetSession = useCallback(() => {
-    setSession({} as Session)
+    setSession({} as InferenceSession)
     setIsLoading(false)
     setError("")
     setIsDisabled(false)
@@ -55,7 +61,7 @@ export function useInferenceSession(
         setQuestion("")
         setError(undefined)
         setIsDisabled(true)
-
+        const messageSentTime = dayjs(Date.now())
         if (globalInference.infer.inferenceId != undefined) {
           setSession((prev) => {
             return {
@@ -88,7 +94,7 @@ export function useInferenceSession(
               signal
             )
 
-            setSession((prev) => {
+            setSession((prev: any) => {
               setStreamMessage("")
               return {
                 ...prev,
@@ -97,10 +103,13 @@ export function useInferenceSession(
                   {
                     from: "agent",
                     text: data,
+                    tokens: data.length,
+                    latency: dayjs(Date.now()).diff(messageSentTime, "seconds"),
                   },
                 ],
               }
             })
+
             setIsLoading(false)
           } catch (e: any) {
             if (e.message === "comet.prompt is not a function") {