AI-Studio/app/MindWork AI Studio/Tools/RAG/RetrievalTextContext.cs

namespace AIStudio.Tools.RAG;

/// <summary>
/// The retrieval context for text data.
/// </summary>
public sealed class RetrievalTextContext : IRetrievalContext
{
    #region Implementation of IRetrievalContext
    
    /// <inheritdoc />
    public required string DataSourceName { get; init; }
    
    /// <inheritdoc />
    public required RetrievalContentCategory Category { get; init; }

    /// <inheritdoc />
    public string SourceCategory { get; init; } = string.Empty;
    
    /// <inheritdoc />
    public required RetrievalContentType Type { get; init; }
    
    /// <inheritdoc />
    public required string Path { get; init; }

    /// <inheritdoc />
    public IReadOnlyList<string> Links { get; init; } = [];

    #endregion

    /// <summary>
    /// The text content.
    /// </summary>
    /// <remarks>
    /// Should contain the matched text and some small context around it.
    /// </remarks>
    public required string MatchedText { get; set; }

    /// <summary>
    /// The surrounding content of the matched text.
    /// </summary>
    /// <remarks>
    /// Might give the user some context about the matched text.
    /// For example, one sentence or paragraph before and after the matched text.
    /// </remarks>
    public IReadOnlyList<string> SurroundingContent { get; set; } = [];

    /// <summary>
    /// Optional title used when this context is displayed as a source reference.
    /// </summary>
    public string ReferenceTitle { get; init; } = string.Empty;

    /// <summary>
    /// Optional link used when this context is displayed as a source reference.
    /// </summary>
    public string ReferenceLink { get; init; } = string.Empty;

    /// <summary>
    /// The page this passage was found on, or null when it has none.
    /// </summary>
    /// <remarks>
    /// Kept as a number rather than only inside the reference title: the AI is told the page so it
    /// can say where an answer comes from, and a source has to name a page a program can be sent
    /// to. A slide or a sheet has no page and leaves this empty.
    /// </remarks>
    public int? PageNumber { get; init; }

    /// <summary>
    /// The page this passage ends on when it goes on beyond the page it starts on, or null when it
    /// stays on that page or the end is not known.
    /// </summary>
    public int? EndPageNumber { get; init; }

    /// <summary>
    /// The sheet of a spreadsheet this passage is from, or null when it is from none.
    /// </summary>
    public string? SheetName { get; init; }

    /// <summary>
    /// The row this passage starts at, counted from one: a row of its sheet, or of a table such as a CSV file.
    /// </summary>
    public int? StartRow { get; init; }

    /// <summary>
    /// The row this passage ends at when it goes on beyond the row it starts at, or null.
    /// </summary>
    public int? EndRow { get; init; }

    /// <summary>
    /// The line this passage starts at, counted from one, for text and source code, or null.
    /// </summary>
    public int? StartLine { get; init; }

    /// <summary>
    /// The line this passage ends at when it goes on beyond the line it starts at, or null.
    /// </summary>
    public int? EndLine { get; init; }

    /// <summary>
    /// Which part of its document this passage is, counted from one, or null when the data source does not say.
    /// </summary>
    /// <remarks>
    /// A document is split into parts to be searched, and each part is one passage. The page says
    /// better where a passage is, so the part only names the place in a document without pages,
    /// such as a text file.
    /// </remarks>
    public int? ChunkNumber { get; init; }

    /// <summary>
    /// Whether AI Studio cut the matched text, since the passage did not fit into a size limit otherwise.
    /// </summary>
    public bool IsMatchedTextShortened { get; private set; }

    /// <summary>
    /// Whether AI Studio cut or left out the surrounding content, since the passage did not fit into a size limit otherwise.
    /// </summary>
    public bool IsSurroundingContentShortened { get; private set; }

    /// <summary>
    /// A copy of this context with shorter texts, made to fit a size limit.
    /// </summary>
    /// <remarks>
    /// Everything else stays as the data source delivered it. The copy remembers which of its texts
    /// are shorter than those of the data source, so that the AI is told.
    /// </remarks>
    /// <param name="matchedText">The matched text, possibly shortened.</param>
    /// <param name="surroundingContent">The surrounding content, possibly shortened.</param>
    /// <returns>The copy.</returns>
    internal RetrievalTextContext WithShorterTexts(string matchedText, IReadOnlyList<string> surroundingContent)
    {
        var copy = (RetrievalTextContext)this.MemberwiseClone();
        copy.IsMatchedTextShortened = this.IsMatchedTextShortened || matchedText.Length < this.MatchedText.Length;
        copy.IsSurroundingContentShortened = this.IsSurroundingContentShortened || !surroundingContent.SequenceEqual(this.SurroundingContent);
        copy.MatchedText = matchedText;
        copy.SurroundingContent = surroundingContent;
        return copy;
    }
}