Merge branch 'main' of github.com:tosoham/StarkFinder

client/app/api/ask/route.ts

@@ -1,10 +1,130 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 // api/ask/route.ts
 import { NextResponse } from 'next/server';
+import { ASK_OPENAI_AGENT_PROMPT } from "@/prompts/prompts";
 import axios from 'axios';
+import { ChatOpenAI } from "@langchain/openai";
+import { ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate } from "@langchain/core/prompts";
+import { START, END, MessagesAnnotation, MemorySaver, StateGraph } from "@langchain/langgraph";
+import { RemoveMessage } from "@langchain/core/messages";
 
 const BRIAN_API_KEY = process.env.BRIAN_API_KEY || '';
-const BRIAN_API_URL = 'https://api.brianknows.org/api/v0/agent';
+const OPENAI_API_KEY = process.env.OPENAI_API_KEY || '';
+const BRIAN_API_URL = 'https://api.brianknows.org/api/v0/agent/knowledge';
+const BRIAN_DEFAULT_RESPONSE: string = "🤖 Sorry, I don’t know how to answer. The AskBrian feature allows you to ask for information on a custom-built knowledge base of resources. Contact the Brian team if you want to add new resources!";
+
+const systemPrompt = ASK_OPENAI_AGENT_PROMPT + `\nThe provided chat history includes a summary of the earlier conversation.`;
+
+const systemMessage = SystemMessagePromptTemplate.fromTemplate([
+  systemPrompt
+]);
+
+const userMessage = HumanMessagePromptTemplate.fromTemplate([
+  "{user_query}"
+]);
+
+const askAgentPromptTemplate = ChatPromptTemplate.fromMessages([
+  systemMessage,
+  userMessage
+]);
+const agent = new ChatOpenAI({
+  modelName: "gpt-4o",
+  temperature: 0.5,
+  openAIApiKey: OPENAI_API_KEY
+});
+const prompt = askAgentPromptTemplate;
+// const chain = prompt.pipe(agent);
+const initialCallModel = async (state: typeof MessagesAnnotation.State) => {
+  const messages = [
+    await systemMessage.format({brianai_answer: BRIAN_DEFAULT_RESPONSE}),
+    ...state.messages
+  ];
+  const response = await agent.invoke(messages);
+  return { messages: response };
+};
+const callModel = async (state: typeof MessagesAnnotation.State ) => {
+  const messageHistory = state.messages.slice(0, -1);
+  if ( messageHistory.length >= 3 ) {
+    const lastHumanMessage = state.messages[state.messages.length - 1];
+    const summaryPrompt = `
+    Distill the above chat messages into a single summary message. 
+    Include as many specific details as you can.
+    IMPORTANT NOTE: Include all information related to user's nature about trading and what kind of trader he/she is. 
+    `;
+    // const summaryMessage = HumanMessagePromptTemplate.fromTemplate([summaryPrompt]);
+    const summary =  await agent.invoke([
+      ...messageHistory,
+      { role: "user", content: summaryPrompt },
+    ]);
+    const deleteMessages = state.messages.map(
+      (m) => m.id ? new RemoveMessage({ id: m.id }) : null
+    );
+    const humanMessage = { role: "user", content: lastHumanMessage.content };
+    const response = await agent.invoke([
+      await systemMessage.format({brianai_answer: BRIAN_DEFAULT_RESPONSE}),
+      summary,
+      humanMessage,
+    ]);
+    //console.log(response);
+    return {
+      messages: [summary, humanMessage, response, ...deleteMessages],
+    };     
+  } else {
+    return await initialCallModel(state);
+  }
+};
+
+const workflow = new StateGraph(MessagesAnnotation)
+  .addNode("model", callModel)
+  .addEdge(START, "model")
+  .addEdge("model", END);
+const app = workflow.compile({ checkpointer: new MemorySaver() });
+
+async function queryOpenAI({userQuery, brianaiResponse}: 
+  {userQuery: string, brianaiResponse: string}): 
+  Promise<string> {
+  try {
+    const response = await app.invoke(
+      {
+        messages: [
+          await prompt.format({brianai_answer: brianaiResponse, user_query: userQuery})
+        ],
+      },
+      {
+        configurable: { thread_id: "1" },
+      },
+    );
+    console.log(response);
+    return response.messages[response.messages.length-1].content as string;
+  } catch (error) {
+    console.error('OpenAI Error:', error);
+    return 'Sorry, I am unable to process your request at the moment.';
+  }
+}
+
+async function queryBrianAI(prompt: string): Promise<string> {
+  try {
+    const response = await axios.post(
+      BRIAN_API_URL,
+      {
+        prompt,
+        kb: "starknet_kb"
+      },
+      {
+        headers: {
+          "Content-Type": "application/json",
+          "x-brian-api-key": BRIAN_API_KEY,
+        }
+      }
+    );
+    const brianaiAnswer = response.data.result.answer;
+    const openaiAnswer = await queryOpenAI({brianaiResponse: brianaiAnswer, userQuery: prompt});
+    return openaiAnswer;
+  } catch (error) {
+    console.error("Brian AI Error:", error);
+    return "Sorry, I am unable to process your request at the moment.";
+  }
+}
 
 export async function POST(request: Request) {
   try {
@@ -33,22 +153,13 @@ export async function POST(request: Request) {
 
     console.log('Request payload:', JSON.stringify(payload, null, 2));
 
-    const response = await axios.post(
-      BRIAN_API_URL,
-      payload,
-      {
-        headers: {
-          'Content-Type': 'application/json',
-          'x-brian-api-key': BRIAN_API_KEY,
-        },
-      }
-    );
+    const response = await queryBrianAI(payload.prompt);
 
-    console.log('API Response:', response.data);
+    console.log('API Response:', response);
 
     // Extract the answer from the result array
-    if (response.data.result && response.data.result[0] && response.data.result[0].answer) {
-      return NextResponse.json({ answer: response.data.result[0].answer });
+    if (response) {
+      return NextResponse.json({ answer: response });
     } else {
       throw new Error('Unexpected API response format');
     }

client/app/api/audit-sourceCode/route.ts

@@ -0,0 +1,132 @@
+import { NextRequest, NextResponse } from 'next/server';
+import { Anthropic } from '@anthropic-ai/sdk';
+
+export async function POST(req: NextRequest) {
+  try {
+    // Parse request body
+    const { sourceCode } = await req.json();
+    const claude = new Anthropic({ apiKey: process.env.ANTHROPIC_API_KEY });
+
+    const stream = await claude.messages.create({
+      model: "claude-3-opus-20240229",
+      system: getStarknetSystemPrompt(),
+      max_tokens: 4096,
+      messages: [
+        {
+          role: "user",
+          content: `Carefully audit the following Starknet smart contract and provide a STRICTLY FORMATTED JSON response:\n\n${sourceCode}`
+        }
+      ],
+      stream: true,
+    });
+
+    const response = new ReadableStream({
+      async start(controller) {
+        let fullResponse = '';
+        for await (const messageStream of stream) {
+          if (messageStream.type === 'content_block_delta') {
+            const deltaText = messageStream.delta.type; // Adjust if incorrect
+            fullResponse += deltaText;
+            controller.enqueue(`data: ${JSON.stringify({ chunk: deltaText })}\n\n`);
+          }
+        }
+        controller.close();
+
+        // Log full response and extract JSON
+        console.log(fullResponse);
+        const jsonContent = extractJSON(fullResponse);
+        try {
+          JSON.parse(jsonContent); // Verify JSON structure
+        } catch (parseError) {
+          console.error('JSON Parsing Error:', parseError);
+          throw new Error('Invalid JSON response received.');
+        }
+      },
+    });
+
+    return new NextResponse(response, {
+      headers: {
+        'Content-Type': 'text/event-stream',
+        'Cache-Control': 'no-cache',
+        Connection: 'keep-alive',
+      },
+    });
+  } catch (error) {
+    console.error('API Error:', error);
+    return NextResponse.json(
+      { error: error instanceof Error ? error.message : 'An unexpected error occurred' },
+      { status: 500 }
+    );
+  }
+}
+
+function extractJSON(text: string) {
+  const codeBlockMatch = text.match(/```json\n([\s\S]*?)```/);
+  if (codeBlockMatch) return codeBlockMatch[1].trim();
+  const bracketMatch = text.match(/\{[\s\S]*\}/);
+  if (bracketMatch) return bracketMatch[0].trim();
+  const cleanedText = text.replace(/^[^{]*/, '').replace(/[^}]*$/, '');
+  return cleanedText;
+}
+
+function getStarknetSystemPrompt() {
+  return `You are a Starknet Smart Contract security expert. Your task is to audit a smart contract focusing on the following security aspects:
+
+1. Contract Anatomy
+- Validate method visibility and access controls
+- Check for proper use of decorators
+- Ensure appropriate function modifiers
+
+2. State Management
+- Verify state mutation safety
+- Check for potential reentrancy vulnerabilities
+- Validate state update patterns
+
+3. Access Control
+- Review authorization mechanisms
+- Check for proper role-based access control
+- Validate ownership and admin privileges
+
+4. External Calls
+- Analyze cross-contract interactions
+- Check for potential manipulation in external calls
+- Verify gas limits and error handling
+
+5. Asset Management
+- Review token transfer mechanisms
+- Check for potential overflow/underflow
+- Validate balance tracking and updates
+
+6. Cryptographic Operations
+- Review signature verification
+- Check for randomness generation
+- Validate cryptographic primitive usage
+
+7. Economic Vulnerabilities
+- Check for potential front-running
+- Analyze economic attack surfaces
+- Verify economic incentive alignment
+
+Output Format:
+{
+    contract_name: string,
+    audit_date: string,
+    security_score: number, // 0-100
+    original_contract_code: string,
+    corrected_contract_code: string,
+    vulnerabilities: [
+        {
+            category: string,
+            severity: 'Low'|'Medium'|'High',
+            description: string,
+            recommended_fix: string
+        }
+    ],
+    recommended_fixes: string[]
+}
+
+IMPORTANT: 
+- Provide the FULL corrected contract code, not just code snippets
+- Include concrete, implementable code fixes for each vulnerability
+- Explain all changes made in the corrected code`;
+}