openai_x_search.go 2.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879
  1. package handler
  2. import (
  3. "encoding/json"
  4. "fmt"
  5. "strings"
  6. "github.com/Wei-Shaw/sub2api/internal/pkg/xai"
  7. "github.com/gin-gonic/gin"
  8. )
  9. type grokStandaloneSearchRequest struct {
  10. Query string `json:"query"`
  11. Input string `json:"input"`
  12. MaxResults *int `json:"max_results"`
  13. AllowedXHandles []string `json:"allowed_x_handles"`
  14. ExcludedXHandles []string `json:"excluded_x_handles"`
  15. FromDate string `json:"from_date"`
  16. ToDate string `json:"to_date"`
  17. EnableImageUnderstanding *bool `json:"enable_image_understanding"`
  18. EnableVideoUnderstanding *bool `json:"enable_video_understanding"`
  19. }
  20. // XSearch marks the standalone endpoint so WebSearch can use native x_search
  21. // while retaining its dedicated per-call billing contract.
  22. func (h *GatewayHandler) XSearch(c *gin.Context) {
  23. c.Set("grok_x_search_endpoint", true)
  24. h.WebSearch(c)
  25. }
  26. func resolveGrokStandaloneSearchModel() string {
  27. return xai.ResolveDefaultTextModel(xai.RuntimeModelMappingOptions().DefaultText)
  28. }
  29. func buildGrokXSearchResponsesBody(req grokStandaloneSearchRequest, model string) ([]byte, error) {
  30. input := strings.TrimSpace(req.Query)
  31. if input == "" {
  32. input = strings.TrimSpace(req.Input)
  33. }
  34. tool := map[string]any{"type": "x_search"}
  35. if len(req.AllowedXHandles) > 0 {
  36. tool["allowed_x_handles"] = req.AllowedXHandles
  37. }
  38. if len(req.ExcludedXHandles) > 0 {
  39. tool["excluded_x_handles"] = req.ExcludedXHandles
  40. }
  41. if strings.TrimSpace(req.FromDate) != "" {
  42. tool["from_date"] = strings.TrimSpace(req.FromDate)
  43. }
  44. if strings.TrimSpace(req.ToDate) != "" {
  45. tool["to_date"] = strings.TrimSpace(req.ToDate)
  46. }
  47. if req.EnableImageUnderstanding != nil {
  48. tool["enable_image_understanding"] = *req.EnableImageUnderstanding
  49. }
  50. if req.EnableVideoUnderstanding != nil {
  51. tool["enable_video_understanding"] = *req.EnableVideoUnderstanding
  52. }
  53. maxResults := 0
  54. if req.MaxResults != nil {
  55. maxResults = *req.MaxResults
  56. }
  57. return json.Marshal(map[string]any{
  58. "model": xai.ResolveDefaultTextModel(model),
  59. "input": buildGrokXSearchPrompt(input, maxResults),
  60. "tools": []map[string]any{tool},
  61. "tool_choice": "required",
  62. "include": []string{"x_search_call.action.sources"},
  63. "store": false,
  64. "stream": false,
  65. })
  66. }
  67. func buildGrokXSearchPrompt(query string, maxResults int) string {
  68. return fmt.Sprintf(`Search X for the user query below. Return ONLY valid JSON with this exact shape: {"results":[{"url":"https://...","title":"post or page title","snippet":"concise factual summary"}]}. Return at most %d unique results. Every URL must be an actual x_search source. Populate a non-empty title and snippet for every result. Do not wrap the JSON in markdown.
  69. User query:
  70. %s`, normalizeGrokWebSearchMaxResults(maxResults), query)
  71. }