MosaicIQ/packages/shared/src/llm/client.ts

/**
 * LLM client for Pi API (Inflection AI)
 * Provides streaming and non-streaming completion functions
 */

const PI_API_URL = "https://api.pi.ai/v1/chat/completions";
const client = {
  apiKey: process.env.PI_API_KEY || "",
};

export interface StreamOptions {
  onProgress?: (text: string) => void;
  signal?: AbortSignal;
}

export interface CompletionOptions {
  model?: string;
  maxTokens?: number;
  temperature?: number;
}

const DEFAULT_MODEL = "pi";
const DEFAULT_MAX_TOKENS = 4096;

/**
 * Stream a response from Pi
 */
export async function streamResponse(
  prompt: string,
  options: StreamOptions & CompletionOptions = {}
): Promise<string> {
  const {
    model = DEFAULT_MODEL,
    maxTokens = DEFAULT_MAX_TOKENS,
    temperature = 0,
    onProgress,
    signal,
  } = options;

  if (!client.apiKey) {
    throw new Error("PI_API_KEY is not set");
  }

  const response = await fetch(PI_API_URL, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      "Authorization": `Bearer ${client.apiKey}`,
    },
    body: JSON.stringify({
      model,
      messages: [{ role: "user", content: prompt }],
      max_tokens: maxTokens,
      temperature,
      stream: true,
    }),
    signal,
  });

  if (!response.ok) {
    const error = await response.text();
    throw new Error(`Pi API error: ${response.status} ${error}`);
  }

  const reader = response.body?.getReader();
  if (!reader) {
    throw new Error("Failed to get response body reader");
  }

  const decoder = new TextDecoder();
  let fullResponse = "";

  while (true) {
    const { done, value } = await reader.read();
    if (done) break;

    const chunk = decoder.decode(value, { stream: true });
    const lines = chunk.split("\n");

    for (const line of lines) {
      if (line.startsWith("data: ") && line !== "data: [DONE]") {
        try {
          const data = JSON.parse(line.slice(6));
          const content = data.choices?.[0]?.delta?.content;
          if (content) {
            fullResponse += content;
            onProgress?.(content);
          }
        } catch {
          // Skip invalid JSON
        }
      }
    }
  }

  return fullResponse;
}

/**
 * Get a complete response from Pi (non-streaming)
 */
export async function complete(
  prompt: string,
  options: CompletionOptions = {}
): Promise<string> {
  const {
    model = DEFAULT_MODEL,
    maxTokens = DEFAULT_MAX_TOKENS,
    temperature = 0,
  } = options;

  if (!client.apiKey) {
    throw new Error("PI_API_KEY is not set");
  }

  const response = await fetch(PI_API_URL, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      "Authorization": `Bearer ${client.apiKey}`,
    },
    body: JSON.stringify({
      model,
      messages: [{ role: "user", content: prompt }],
      max_tokens: maxTokens,
      temperature,
    }),
  });

  if (!response.ok) {
    const error = await response.text();
    throw new Error(`Pi API error: ${response.status} ${error}`);
  }

  const data = await response.json();
  return data.choices?.[0]?.message?.content ?? "";
}

/**
 * Check if API key is configured
 */
export function isConfigured(): boolean {
  return !!client.apiKey && client.apiKey.length > 0;
}

/**
 * Stream structured JSON response
 * Useful for agents that need to return structured data
 */
export async function streamStructuredResponse<T>(
  prompt: string,
  schema: Record<string, unknown>,
  options: StreamOptions & CompletionOptions = {}
): Promise<T> {
  const structuredPrompt = `${prompt}

Please respond with a JSON object that follows this structure:
${JSON.stringify(schema, null, 2)}

Your response must be valid JSON only, with no additional text or explanation.`;

  const response = await streamResponse(structuredPrompt, options);

  // Try to parse as JSON
  try {
    return JSON.parse(response) as T;
  } catch (error) {
    // If the response isn't valid JSON, try to extract JSON from the response
    const jsonMatch = response.match(/\{[\s\S]*\}/);
    if (jsonMatch) {
      return JSON.parse(jsonMatch[0]) as T;
    }
    throw new Error("Failed to parse structured response as JSON");
  }
}

/**
 * Complete with structured JSON response (non-streaming)
 */
export async function completeStructured<T>(
  prompt: string,
  schema: Record<string, unknown>,
  options: CompletionOptions = {}
): Promise<T> {
  const structuredPrompt = `${prompt}

Please respond with a JSON object that follows this structure:
${JSON.stringify(schema, null, 2)}

Your response must be valid JSON only, with no additional text or explanation.`;

  const response = await complete(structuredPrompt, options);

  // Try to parse as JSON
  try {
    return JSON.parse(response) as T;
  } catch (error) {
    // If the response isn't valid JSON, try to extract JSON from the response
    const jsonMatch = response.match(/\{[\s\S]*\}/);
    if (jsonMatch) {
      return JSON.parse(jsonMatch[0]) as T;
    }
    throw new Error("Failed to parse structured response as JSON");
  }
}

export default client;