add ability to ask question and get related files using embeddings wi…

…th ai answer.

package.json

@@ -21,6 +21,7 @@
     "typecheck": "tsc --noEmit"
   },
   "dependencies": {
+    "@ai-sdk/google": "^1.0.4",
     "@clerk/nextjs": "^6.5.0",
     "@google/generative-ai": "^0.21.0",
     "@hookform/resolvers": "^3.9.1",
@@ -59,6 +60,9 @@
     "@trpc/client": "^11.0.0-rc.446",
     "@trpc/react-query": "^11.0.0-rc.446",
     "@trpc/server": "^11.0.0-rc.446",
+    "@types/react-syntax-highlighter": "^15.5.13",
+    "@uiw/react-md-editor": "^4.0.4",
+    "ai": "^4.0.11",
     "axios": "^1.7.8",
     "class-variance-authority": "^0.7.1",
     "clsx": "^2.1.1",
@@ -77,6 +81,7 @@
     "react-dom": "^18.3.1",
     "react-hook-form": "^7.53.2",
     "react-resizable-panels": "^2.1.7",
+    "react-syntax-highlighter": "^15.6.1",
     "recharts": "^2.13.3",
     "server-only": "^0.0.1",
     "shadcn-ui": "^0.9.3",

prisma/schema.prisma

@@ -25,6 +25,7 @@ model User {
 
     credits Int @default(150)
     userToProjects UserToProject[]
+    questionsAsked Question[]
 }
 
 model Project {
@@ -39,6 +40,8 @@ model Project {
     userToProjects UserToProject[]
     commits Commit[]
     sourceCodeEmbeddings SourceCodeEmbedding[]
+
+    savedQuestions Question[]
 }
 
 model UserToProject {
@@ -67,6 +70,23 @@ model SourceCodeEmbedding {
     project Project @relation(fields: [projectId], references: [id])
 }
 
+model Question {
+    id      String  @id @default(cuid())
+    createdAt DateTime @default(now())
+    updatedAt DateTime @updatedAt
+
+    question String
+    answer String
+
+    fileReferences Json?
+
+    projectId String
+    project Project @relation(fields: [projectId], references: [id])
+
+    userId String
+    user User @relation(fields: [userId], references: [id])
+}
+
 model Commit {
     id      String  @id @default(cuid())
     createdAt DateTime @default(now())

src/app/(protected)/dashboard/actions.ts

@@ -0,0 +1,63 @@
+'use server'
+import {streamText} from "ai";
+import {createStreamableValue} from 'ai/rsc'
+import {createGoogleGenerativeAI} from '@ai-sdk/google'
+import { generateEmbedding } from "@/lib/gemini";
+import { db } from "@/server/db";
+
+const google = createGoogleGenerativeAI({
+    apiKey: process.env.GEMINI_API_KEY,
+})
+
+export async function askQuestion(question: string, projectId: string) {
+    const stream = createStreamableValue()
+
+    const queryVector = await generateEmbedding(question)
+    const vectorQuery = `[[${queryVector.join(',')}]`
+    const result = await db.$queryRaw`SELECT "fileName", "sourceCode", "summary", 1-("summaryEmbedding" <=> ${vectorQuery}::vector) AS similarity FROM "sourceCodeEmbedding" WHERE 1-("summaryEmbedding" <=> ${vectorQuery}::vector) > .5 AND "projectId" = ${projectId} ORDER BY similarity DESC LIMIT 10` as {fileName: string, sourceCode: string, summary: string}[]
+
+    let context = ''
+    for(const doc of result){
+        context += `source: ${doc.fileName}\ncode content: ${doc.sourceCode}\n summary of file: ${doc.summary}\n\n`
+    }
+
+    (async ()=> {
+        const {textStream} = await streamText({
+            model: google('gemini-1.5-flash'),
+            prompt: `
+            You are an ai code assistant who answers questions about the codebase. Your target audience is a technical intern.
+            AI assistant is a brand new, powerful, human-like artificial intellgience.
+            The traits of AI include expert knowledge, helpfulness, cleverness, and articulateness.
+            AI is well-behaved and well-mannered individual.
+            AI is always friendly, kind and inspiring, and he is eager to provide vivid and thoughtful responses to the user.
+            AI has the sum of all knowledge in their brain, and is able to accurately answer nearly any question about any topic.
+            If the question is asking about code or a specific file, AI will provide the detailed answer, giving step by step instructions.
+            START CONTEXT BLOCK
+            ${context}
+            END CONTEXT BLOCK
+
+            START QUESTION
+            ${question}
+            END QUESTION
+            AI Assistant will take into account any CONTEXT BLOCK that is provided in a conversation.
+            If the context does not provide the answer to question, AI Assistant will say, "I'm sorry, but I don't know the answer."
+            AI assistant will not apologize for previous responses, but instead will indicate new information was gained.
+            AI Assistant will not invent anything that is not drawn directly from the context.
+            Answer in markdown syntax, with code snippets if needed. Be as detailed as possible when answering, make sure there is no missing information.
+            `
+        });
+
+        for await (const delta of textStream) {
+            stream.update(delta)
+        }
+
+        stream.done()
+
+    })()
+
+    return {
+        output: stream.value,
+        fileReferences: result
+    }
+
+}

src/app/(protected)/dashboard/ask-question-card.tsx

@@ -0,0 +1,96 @@
+'use client'
+import MDEditor from '@uiw/react-md-editor'
+import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card'
+import { Textarea } from '@/components/ui/textarea'
+import useProject from '@/hooks/use-project'
+import React from 'react'
+import { Button } from '@/components/ui/button';
+import { Dialog, DialogContent, DialogHeader, DialogTitle } from '@/components/ui/dialog'
+import Image from 'next/image'
+import { askQuestion } from './actions'
+import { string } from 'zod'
+import { readStreamableValue } from 'ai/rsc'
+import CodeReferences from './code-references'
+import { api } from '@/trpc/react'
+import { toast } from 'sonner'
+
+const saveAnswer = api.project.saveAnswer.useMutation()
+
+const AskQuestionCard = () => {
+    const {project} = useProject()
+    const [open, setOpen] = React.useState(false)
+    const [question, setQuestion] = React.useState('')
+    const [loading, setLoading] = React.useState(false)
+    const [fileReferences, setFileReferences] = React.useState<{fileName: string; sourceCode: string; summary: string}[]>([])
+    const [answer, setAnswer] = React.useState('')
+
+    const onSubmit = async (e: React.FormEvent<HTMLFormElement>) => {
+        setAnswer('')
+    setFileReferences([])
+        e.preventDefault()
+        if(!project?.id) return
+        setLoading(true)
+
+
+        const {output, fileReferences} = await askQuestion(question, project.id)
+        setOpen(true)
+        setFileReferences(fileReferences)
+
+        for await (const delta of readStreamableValue(output)) {
+            if(delta){
+                setAnswer((ans) => ans + delta)
+            }
+        }
+        setLoading(false)
+    }
+
+  return (
+    <>
+    <Dialog open={open} onOpenChange={setOpen}>
+        <DialogContent className="sm:max-w-[80vw]">
+        <DialogHeader>
+            <div className="flex items-center gap-2">
+            </div>
+            <DialogTitle>
+                <Image src="/logo.png" width={40} height={40} alt="gitSage" />
+            </DialogTitle>
+            <Button disabled={saveAnswer.isPending} variant="outline" onClick={() => {saveAnswer.mutate({projectId: project!.id, question, answer, fileReferences}, {
+                onSuccess: () => {
+                    toast.success('Answer saved successfully')
+                },
+                onError: () => {
+                    toast.error('Error saving answer')
+                }
+            })}}>
+                Save Answer
+            </Button>
+        </DialogHeader>
+        <MDEditor.Markdown source={answer} className="max-w-[70vw] !h-full max-h-[40vh] overflow-scroll"/>
+        <div className="h-4"></div>
+        <CodeReferences fileReferences={fileReferences} />
+
+        <Button type="button" onClick={() => {setOpen(false)}}>
+        Close
+        </Button>
+        </DialogContent>
+
+    </Dialog>
+    <Card className='relative col-span-3 '>
+        <CardHeader>
+            <CardTitle>Ask a question</CardTitle>
+        </CardHeader>
+        <CardContent>
+            <form onSubmit={onSubmit}>
+                <Textarea placeholder='Which file should I edit to change the home page?' />
+                <div className="h-4"></div>
+                <Button type="submit" disabled={loading}>
+                    Ask GitSage!
+                </Button>
+            </form>
+        </CardContent>
+    </Card>
+    </>
+  )
+}
+
+export default AskQuestionCard

src/app/(protected)/dashboard/code-references.tsx

@@ -0,0 +1,44 @@
+'use client'
+
+import React from 'react'
+import {Tabs, TabsContent, TabsList, TabsTrigger} from "@/components/ui/tabs"
+import {Prism as SyntaxHighlighter} from 'react-syntax-highlighter'
+import {dracula} from 'react-syntax-highlighter/dist/esm/styles/prism'
+import { cn } from '@/lib/utils'
+
+type Props = {
+    fileReferences: { fileName: string; sourceCode: string; summary: string}[]
+}
+
+const CodeReferences = ({fileReferences}: Props) => {
+    const [tab, setTab] = React.useState(fileReferences[0]?.fileName)
+    if(fileReferences.length === 0) return null
+
+  return (
+    <div className='max-w-[70vw]'>
+        <Tabs value={tab} onValueChange={setTab}>
+            <div className='overflow-scroll flex gap-2 bg-gray-200 p-1 rounded-md'>
+                {fileReferences.map((file) => (
+                    <button onClick={() => setTab(file.fileName)} key={file.fileName} className={cn(
+                        'px-3 py-1.5 text-sm font-medium rounded-md transition-colors whitespace-nowrap text-muted-foreground hover:bg-muted',
+                        {
+                            'bg-primary text-primary-foreground': tab === file.fileName,
+                        }
+                    )}>
+                        {file.fileName}
+                    </button>
+                ))}
+            </div>
+            {fileReferences.map(file => (
+                <TabsContent key={file.fileName} value={file.fileName} className='max-h-[40vh] overflow-scroll max-w-7xl rounded-md'>
+                    <SyntaxHighlighter language="typescript" style={dracula}>
+                        {file.sourceCode}
+                    </SyntaxHighlighter>
+                </TabsContent>
+            ))}
+        </Tabs>
+    </div>
+  )
+}
+
+export default CodeReferences

src/app/(protected)/dashboard/page.tsx

@@ -6,6 +6,7 @@ import { ExternalLink, Github } from "lucide-react";
 import Link from "next/link";
 import React from "react";
 import CommitLog from "./commit-log";
+import AskQuestionCard from "./ask-question-card";
 
 const DashboardPage = () => {
   const {project} = useProject()
@@ -39,7 +40,7 @@ const DashboardPage = () => {
 
       <div className="mt-4">
         <div className="grid grid-cols-1 gap-4 sm:grid-cols-5">
-          AskQuestionCard
+          <AskQuestionCard />
           MeetingCard
         </div>
       </div>

src/server/api/routers/project.ts

@@ -49,6 +49,21 @@ export const projectRouter = createTRPCRouter({
                 projectId: input.projectId
             }
         })
+    }),
+    saveAnswer: protectedProcedure.input(z.object({
+        projectId: z.string(),
+        question: z.string(),
+        answer: z.string(),
+        fileReferences: z.any()
+    })).mutation(async ({ ctx, input }) => {
+        return await ctx.db.question.create({
+            data: {
+                answer: input.answer,
+                fileReferences: input.fileReferences,
+                projectId: input.projectId,
+                question: input.question,
+                userId: ctx.user.userId!
+            }
+        })
     })
-
 });