File size: 4,817 Bytes
93c19dc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
import { COOKIE_NAME } from "@shared/const";
import { getSessionCookieOptions } from "./_core/cookies";
import { systemRouter } from "./_core/systemRouter";
import { publicProcedure, protectedProcedure, router } from "./_core/trpc";
import { z } from "zod";
import { generateResponseWithReasoning, generateImage } from "./llm";
import { searchOnline, formatSearchResults, sanitizeSearchQuery } from "./search";
import { checkRateLimit } from "./rateLimit";

export const appRouter = router({
  // if you need to use socket.io, read and register route in server/_core/index.ts, all api should start with '/api/' so that the gateway can route correctly
  system: systemRouter,
  auth: router({
    me: publicProcedure.query(opts => opts.ctx.user),
    logout: publicProcedure.mutation(({ ctx }) => {
      const cookieOptions = getSessionCookieOptions(ctx.req);
      ctx.res.clearCookie(COOKIE_NAME, { ...cookieOptions, maxAge: -1 });
      return {
        success: true,
      } as const;
    }),
  }),

  // Section 1: Chat procedures (Ask mode)
  chat: router({
    send: protectedProcedure
      .input(
        z.object({
          prompt: z.string().min(1),
          conversationId: z.number().optional(),
          enableSearch: z.boolean().default(false),
          enableThinking: z.boolean().default(false),
          history: z
            .array(
              z.object({
                role: z.enum(["user", "assistant"]),
                content: z.string(),
              })
            )
            .default([]),
        })
      )
      .mutation(async ({ ctx, input }) => {
        // Check rate limit
        const userId = ctx.user?.id?.toString() || "anonymous";
        const rateLimitCheck = checkRateLimit(userId, "chat");

        if (!rateLimitCheck.allowed) {
          throw new Error(
            `Rate limit exceeded. Try again in ${Math.ceil(rateLimitCheck.resetIn / 1000)} seconds.`
          );
        }

        try {
          let searchResults = "";

          // Search online if enabled
          if (input.enableSearch) {
            const sanitizedQuery = sanitizeSearchQuery(input.prompt);
            const results = await searchOnline(sanitizedQuery, 5);
            searchResults = formatSearchResults(results);
          }

          // Generate response with optional reasoning
          const response = await generateResponseWithReasoning(
            input.prompt,
            searchResults,
            input.enableThinking,
            input.history
          );

          return {
            success: true,
            response: response.response,
            reasoning: response.reasoning,
            model: response.model,
            tokensUsed: response.tokensUsed,
            searchResults: searchResults || undefined,
          };
        } catch (error) {
          console.error("Chat error:", error);
          throw new Error(
            error instanceof Error ? error.message : "Failed to generate response"
          );
        }
      }),
  }),

  // Section 1: Image generation procedures (Imagine mode)
  imagine: router({
    generate: protectedProcedure
      .input(
        z.object({
          prompt: z.string().min(1),
        })
      )
      .mutation(async ({ ctx, input }) => {
        // Check rate limit
        const userId = ctx.user?.id?.toString() || "anonymous";
        const rateLimitCheck = checkRateLimit(userId, "imagine");

        if (!rateLimitCheck.allowed) {
          throw new Error(
            `Rate limit exceeded. Try again in ${Math.ceil(rateLimitCheck.resetIn / 1000)} seconds.`
          );
        }

        try {
          const imageUrl = await generateImage(input.prompt);
          return {
            success: true,
            imageUrl,
            prompt: input.prompt,
          };
        } catch (error) {
          console.error("Image generation error:", error);
          throw new Error(
            error instanceof Error ? error.message : "Failed to generate image"
          );
        }
      }),
  }),

  // Section 1: Search procedure
  search: router({
    online: publicProcedure
      .input(
        z.object({
          query: z.string().min(1),
          maxResults: z.number().default(5),
        })
      )
      .query(async ({ input }) => {
        try {
          const sanitizedQuery = sanitizeSearchQuery(input.query);
          const results = await searchOnline(sanitizedQuery, input.maxResults);
          return {
            success: true,
            results,
            query: input.query,
          };
        } catch (error) {
          console.error("Search error:", error);
          throw new Error(
            error instanceof Error ? error.message : "Search failed"
          );
        }
      }),
  }),
});

export type AppRouter = typeof appRouter;