ai-robot-core/ai-service-admin/src/api/rag.ts

import request from '@/utils/request'

export interface AIResponse {
  content: string
  prompt_tokens?: number
  completion_tokens?: number
  total_tokens?: number
  latency_ms?: number
  model?: string
}

export interface RetrievalResult {
  content: string
  score: number
  source: string
  metadata?: Record<string, any>
}

export interface RagExperimentRequest {
  query: string
  kb_ids?: string[]
  top_k?: number
  score_threshold?: number
  llm_provider?: string
  generate_response?: boolean
}

export interface RagExperimentResult {
  query: string
  retrieval_results?: RetrievalResult[]
  final_prompt?: string
  ai_response?: AIResponse
  total_latency_ms?: number
}

export function runRagExperiment(data: RagExperimentRequest): Promise<RagExperimentResult> {
  return request({
    url: '/admin/rag/experiments/run',
    method: 'post',
    data
  })
}

export function runRagExperimentStream(
  data: RagExperimentRequest,
  onMessage: (event: MessageEvent) => void,
  onError?: (error: Event) => void,
  onComplete?: () => void
): EventSource {
  const baseUrl = import.meta.env.VITE_APP_BASE_API || '/api'
  const url = `${baseUrl}/admin/rag/experiments/stream`
  
  const eventSource = new EventSource(url, {
    withCredentials: true
  })
  
  eventSource.onmessage = onMessage
  eventSource.onerror = (error) => {
    eventSource.close()
    onError?.(error)
  }
  
  return eventSource
}

export function createSSEConnection(
  url: string,
  body: RagExperimentRequest,
  onMessage: (data: string) => void,
  onError?: (error: Error) => void,
  onComplete?: () => void
): () => void {
  const baseUrl = import.meta.env.VITE_APP_BASE_API || '/api'
  const fullUrl = `${baseUrl}${url}`
  
  const controller = new AbortController()
  
  fetch(fullUrl, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      'Accept': 'text/event-stream',
    },
    body: JSON.stringify(body),
    signal: controller.signal
  })
    .then(async (response) => {
      if (!response.ok) {
        throw new Error(`HTTP error! status: ${response.status}`)
      }
      
      const reader = response.body?.getReader()
      if (!reader) {
        throw new Error('No response body')
      }
      
      const decoder = new TextDecoder()
      let buffer = ''
      
      while (true) {
        const { done, value } = await reader.read()
        
        if (done) {
          onComplete?.()
          break
        }
        
        buffer += decoder.decode(value, { stream: true })
        const lines = buffer.split('\n')
        buffer = lines.pop() || ''
        
        for (const line of lines) {
          if (line.startsWith('data: ')) {
            const data = line.slice(6)
            if (data === '[DONE]') {
              onComplete?.()
              return
            }
            onMessage(data)
          }
        }
      }
    })
    .catch((error) => {
      if (error.name !== 'AbortError') {
        onError?.(error)
      }
    })
  
  return () => controller.abort()
}