Generate textbook-quality LLM pretraining data
import asyncio
import os
from logging.config import fileConfig

from alembic import context
from sqlalchemy import engine_from_config
from sqlalchemy import pool
from sqlalchemy.ext.asyncio import async_engine_from_config
from sqlmodel import SQLModel


# this is the Alembic Config object, which provides
# access to the values within the .ini file in use.
config = context.config

# Interpret the config file for Python logging.
# This line sets up loggers basically.
if config.config_file_name is not None:
    fileConfig(config.config_file_name)

# add your model's MetaData object here
# for 'autogenerate' support

from app.db.tables import BaseDBModel

target_metadata = SQLModel.metadata


# other values from the config, defined by the needs of env.py,
# can be acquired:
# my_important_option = config.get_main_option("my_important_option")
# ... etc.


def get_url():
    conn = os.getenv("DATABASE_URL")
    if conn and conn.startswith("postgres://"):
        conn = conn.replace("postgres://", "postgresql+asyncpg://", 1)
    elif conn and conn.startswith("postgresql://"):
        conn = conn.replace("postgresql://", "postgresql+asyncpg://", 1)
    return conn


def do_run_migrations(connection):
    context.configure(connection=connection, target_metadata=target_metadata)

    with context.begin_transaction():
        context.run_migrations()


async def run_async_migrations():
    """In this scenario we need to create an Engine
    and associate a connection with the context.

    """
    configuration = config.get_section(config.config_ini_section)
    configuration["sqlalchemy.url"] = get_url()

    connectable = async_engine_from_config(
        configuration,
        prefix="sqlalchemy.",
        poolclass=pool.NullPool,
    )

    async with connectable.connect() as connection:
        await connection.run_sync(do_run_migrations)

    await connectable.dispose()


def run_migrations_offline() -> None:
    """Run migrations in 'offline' mode.

    This configures the context with just a URL
    and not an Engine, though an Engine is acceptable
    here as well.  By skipping the Engine creation
    we don't even need a DBAPI to be available.

    Calls to context.execute() here emit the given string to the
    script output.

    """
    url = get_url()
    context.configure(
        url=url,
        target_metadata=target_metadata,
        literal_binds=True,
        dialect_opts={"paramstyle": "named"},
    )

    with context.begin_transaction():
        context.run_migrations()


def run_migrations_online() -> None:
    """Run migrations in 'online' mode.

    In this scenario we need to create an Engine
    and associate a connection with the context.

    """
    asyncio.run(run_async_migrations())


if context.is_offline_mode():
    run_migrations_offline()
else:
    run_migrations_online()

alembic/env.py
VikParuchuri/textbook_quality
"""empty message

Revision ID: 5032f4520049
Revises: 
Create Date: 2023-09-20 12:19:11.011545

"""
from alembic import op
import sqlalchemy as sa
import sqlmodel
import app


# revision identifiers, used by Alembic.
revision = '5032f4520049'
down_revision = None
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('prompt',
                    sa.Column('created', app.db.base_model.TZDateTime(timezone=True), nullable=True),
                    sa.Column('updated', app.db.base_model.TZDateTime(timezone=True), nullable=True),
                    sa.Column('id', sa.Integer(), nullable=False),
                    sa.Column('hash', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
                    sa.Column('prompt', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
                    sa.Column('response', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
                    sa.Column('type', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
                    sa.Column('model', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
                    sa.PrimaryKeyConstraint('id'),
                    sa.UniqueConstraint('hash', 'model', name='unique_hash_model')
                    )
    op.create_index(op.f('ix_prompt_hash'), 'prompt', ['hash'], unique=False)
    op.create_index(op.f('ix_prompt_id'), 'prompt', ['id'], unique=False)
    op.create_table('scrapeddata',
                    sa.Column('created', app.db.base_model.TZDateTime(timezone=True), nullable=True),
                    sa.Column('updated', app.db.base_model.TZDateTime(timezone=True), nullable=True),
                    sa.Column('extra', sa.JSON(), nullable=True),
                    sa.Column('id', sa.Integer(), nullable=False),
                    sa.Column('source', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
                    sa.Column('uploaded', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
                    sa.PrimaryKeyConstraint('id')
                    )
    op.create_index(op.f('ix_scrapeddata_id'), 'scrapeddata', ['id'], unique=False)
    op.create_index(op.f('ix_scrapeddata_source'), 'scrapeddata', ['source'], unique=True)
    op.create_table('serviceresponse',
                    sa.Column('created', app.db.base_model.TZDateTime(timezone=True), nullable=True),
                    sa.Column('updated', app.db.base_model.TZDateTime(timezone=True), nullable=True),
                    sa.Column('request', sa.JSON(), nullable=True),
                    sa.Column('response', sa.JSON(), nullable=True),
                    sa.Column('id', sa.Integer(), nullable=False),
                    sa.Column('hash', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
                    sa.Column('name', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
                    sa.PrimaryKeyConstraint('id'),
                    sa.UniqueConstraint('hash', 'name', name='unique_hash_name')
                    )
    op.create_index(op.f('ix_serviceresponse_hash'), 'serviceresponse', ['hash'], unique=False)
    op.create_index(op.f('ix_serviceresponse_id'), 'serviceresponse', ['id'], unique=False)
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(op.f('ix_serviceresponse_id'), table_name='serviceresponse')
    op.drop_index(op.f('ix_serviceresponse_hash'), table_name='serviceresponse')
    op.drop_table('serviceresponse')
    op.drop_index(op.f('ix_scrapeddata_source'), table_name='scrapeddata')
    op.drop_index(op.f('ix_scrapeddata_id'), table_name='scrapeddata')
    op.drop_table('scrapeddata')
    op.drop_index(op.f('ix_prompt_id'), table_name='prompt')
    op.drop_index(op.f('ix_prompt_hash'), table_name='prompt')
    op.drop_table('prompt')
    # ### end Alembic commands ###

alembic/versions/2023_09_20_1219-5032f4520049_.py
VikParuchuri/textbook_quality
"""empty message

Revision ID: 319c7bd1d8b7
Revises: 5032f4520049
Create Date: 2023-09-26 11:56:22.821284

"""
from alembic import op
import sqlalchemy as sa
import sqlmodel
import app


# revision identifiers, used by Alembic.
revision = '319c7bd1d8b7'
down_revision = '5032f4520049'
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('course',
                    sa.Column('created', app.db.base_model.TZDateTime(timezone=True), nullable=True),
                    sa.Column('updated', app.db.base_model.TZDateTime(timezone=True), nullable=True),
                    sa.Column('outline', sa.JSON(), nullable=True),
                    sa.Column('concepts', sa.JSON(), nullable=True),
                    sa.Column('markdown', sa.JSON(), nullable=True),
                    sa.Column('components', sa.JSON(), nullable=True),
                    sa.Column('context', sa.JSON(), nullable=True),
                    sa.Column('id', sa.Integer(), nullable=False),
                    sa.Column('model', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
                    sa.Column('topic', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
                    sa.PrimaryKeyConstraint('id'),
                    sa.UniqueConstraint('topic', 'model', name='unique_topic_model')
                    )
    op.create_index(op.f('ix_course_id'), 'course', ['id'], unique=False)
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(op.f('ix_course_id'), table_name='course')
    op.drop_table('course')
    # ### end Alembic commands ###

alembic/versions/2023_09_26_1156-319c7bd1d8b7_.py
VikParuchuri/textbook_quality
"""empty message

Revision ID: 45f5c62a8b98
Revises: 319c7bd1d8b7
Create Date: 2023-10-02 14:25:16.740000

"""
from alembic import op
import sqlalchemy as sa
import sqlmodel
import app


# revision identifiers, used by Alembic.
revision = '45f5c62a8b98'
down_revision = '319c7bd1d8b7'
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('course', sa.Column('queries', sa.JSON(), nullable=True))
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('course', 'queries')
    # ### end Alembic commands ###

alembic/versions/2023_10_02_1425-45f5c62a8b98_.py
VikParuchuri/textbook_quality
"""empty message

Revision ID: dcff5f57b3b6
Revises: 45f5c62a8b98
Create Date: 2023-10-05 18:20:22.249295

"""
from alembic import op
import sqlalchemy as sa
import sqlmodel
import app


# revision identifiers, used by Alembic.
revision = 'dcff5f57b3b6'
down_revision = '45f5c62a8b98'
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('course', sa.Column('version', sa.Integer(), nullable=False, server_default="1"))
    op.drop_constraint('unique_topic_model', 'course', type_='unique')
    op.create_unique_constraint('unique_topic_model_version', 'course', ['topic', 'model', 'version'])
    op.add_column('prompt', sa.Column('version', sa.Integer(), nullable=False, server_default="1"))
    op.drop_constraint('unique_hash_model', 'prompt', type_='unique')
    op.create_unique_constraint('unique_hash_model_version', 'prompt', ['hash', 'model', 'version'])
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint('unique_hash_model_version', 'prompt', type_='unique')
    op.create_unique_constraint('unique_hash_model', 'prompt', ['hash', 'model'])
    op.drop_column('prompt', 'version')
    op.drop_constraint('unique_topic_model_version', 'course', type_='unique')
    op.create_unique_constraint('unique_topic_model', 'course', ['topic', 'model'])
    op.drop_column('course', 'version')
    # ### end Alembic commands ###

alembic/versions/2023_10_05_1820-dcff5f57b3b6_.py
VikParuchuri/textbook_quality
import json
import re
import secrets
from typing import get_args, List, Dict

from app.components.schemas import (
    COMPONENT_MAP,
    COMPONENT_VERSION,
    ComponentNames,
    ExerciseModes, AllLessonComponentData,
)
from app.settings import settings


def parse_components_from_dicts(components: List[Dict]) -> List[AllLessonComponentData]:
    new_components = []
    for component in components:
        component_cls = COMPONENT_MAP[component["type"]]
        new_components.append(component_cls(**component))
    return new_components


def parse_single_component(component_type: str, markdown: str, key: str = None):
    component_cls = COMPONENT_MAP[component_type]
    component_data = {
        "key": key if key else generate_component_key(),
        "version": COMPONENT_VERSION,
        "type": component_type,
    }
    match component_type:
        case ComponentNames.exercise:
            exercise_data = re.split(r"\nAnswer\n|\nAnswer:\n|\nAnswer: |\nSolution: |\nSolution:\n|\nSolution\n", markdown)
            match len(exercise_data):
                case 1:
                    instructions, answer = exercise_data[0], None
                case 2:
                    instructions, answer = exercise_data
                    answer = answer.strip()
                case _:
                    instructions = exercise_data[0]
                    answer = "\n".join(exercise_data[1:])
            instructions = instructions.strip()
            instructions = (
                instructions.replace("Instructions\n", "").
                replace("Instructions:\n", "")
                .replace("\nExercise\n", "")
                .replace("\nExercise:\n", "")
                .strip()
            )

            mode = (
                ExerciseModes.code
                if answer and "```" in answer
                else ExerciseModes.short_answer
            )

            new_component = component_cls(
                markdown=markdown,
                instructions=instructions,
                solution=answer,
                mode=mode,
                **component_data,
            )
        case ComponentNames.section:
            new_component = component_cls(markdown=markdown, **component_data)
        case _:
            if component_type not in list(
                get_args(settings.VALID_GENERATED_COMPONENTS)
            ):
                raise ValueError(f"Invalid component type: {component_type}")
            new_component = component_cls(markdown=markdown, **component_data)
    return new_component


def generate_component_key() -> str:
    return f"ek-{secrets.token_hex(16)}"

app/components/parser.py
VikParuchuri/textbook_quality
from typing import Literal, Optional, Union

import markdown
from pydantic import BaseModel, validator

from app.settings import settings
from app.util import BaseEnum


class LessonComponentData(BaseModel):
    key: str
    type: settings.VALID_COMPONENTS  # these are the component types
    markdown: str = ""
    version: int = 1

class MarkdownComponentData(LessonComponentData):
    type: Literal["text", "example", "diagram", "section"]

    @validator("markdown")
    def validate_markdown(cls, value):
        try:
            # Try to parse the text as markdown to html
            markdown.markdown(value)
        except Exception:
            # If parsing fails, raise a validation error
            raise ValueError("Text must be valid markdown")
        return value


class ExerciseModes(str, BaseEnum):
    short_answer = "short_answer"
    code = "code"


class ExerciseComponentData(LessonComponentData):
    type: Literal["exercise"] = "exercise"
    markdown: Optional[str]
    instructions: Optional[str]
    solution: Optional[str]
    mode: Optional[ExerciseModes] = ExerciseModes.short_answer


AllLessonComponentData = Union[
    ExerciseComponentData,
    MarkdownComponentData,
]


class ComponentNames(str, BaseEnum):
    text = "text"
    example = "example"
    exercise = "exercise"
    section = "section"


COMPONENT_MAP = {
    ComponentNames.text: MarkdownComponentData,
    ComponentNames.example: MarkdownComponentData,
    ComponentNames.exercise: ExerciseComponentData,
    ComponentNames.section: MarkdownComponentData,
}

COMPONENT_VERSION = 1

app/components/schemas.py
VikParuchuri/textbook_quality
import os
from typing import List

import torch
from sentence_transformers import util, SentenceTransformer

from app.course.schemas import ResearchNote

EMBEDDING_DIM = 384

os.environ["TOKENIZERS_PARALLELISM"] = "false"


def create_embeddings(passages, model) -> torch.Tensor:
    return model.encode(passages, convert_to_tensor=True)


def run_query(
    query_text: str | List[str], embeddings, model, result_count=1, score_thresh=0.6
):
    query_embedding = model.encode(query_text, convert_to_tensor=True)

    cos_scores = util.cos_sim(query_embedding, embeddings)
    top_results = torch.topk(cos_scores, k=result_count, dim=-1)

    # Indices of the passages most similar to the queries (outline items)
    flat_indices = torch.flatten(top_results.indices[top_results.values > score_thresh])
    selected_row = torch.sum(top_results.values > score_thresh, dim=-1) > 0
    selected_indices = set(flat_indices.tolist())

    # Create a mapping, so we know which passages are used by which outline items
    if isinstance(query_text, list):
        outline_items_selected = torch.arange(len(query_text))[selected_row].tolist()
        item_mapping = {o: i.item() for o, i in zip(outline_items_selected, flat_indices)}
    else:
        item_mapping = None
    return top_results, selected_indices, item_mapping


def dedup_list(topics, score_thresh=0.9):
    tc = TopicEmbedding()
    clean_topics = topics[:1]
    tc.add_topics(clean_topics)

    for topic in topics[1:]:
        res = tc.query(topic, score_thresh=score_thresh)
        if len(res) == 0:
            clean_topics.append(topic)
        try:
            tc.add_topics([topic])
        except KeyError:
            pass
    return clean_topics


class TopicEmbedding:
    def __init__(self):
        self.embeddings = None
        self.topics = []
        self.model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

    def add_topics(self, topics):
        for topic in topics:
            self.topics.append(topic)
            embeddings = create_embeddings(topic, self.model)

            embeddings = embeddings.reshape(1, -1)
            if embeddings.shape[1] != EMBEDDING_DIM:
                print(f"Error embedding topic: {topic}")
                continue

            if self.embeddings is None:
                self.embeddings = embeddings
            else:
                self.embeddings = torch.cat((self.embeddings, embeddings), dim=0)

    def query(self, query_text, result_count=1, score_thresh=0.9) -> List[str]:
        try:
            scores, selected_indices, item_mapping = run_query(
                query_text, self.embeddings, self.model, result_count, score_thresh=score_thresh
            )
        except KeyError as e:
            print(f"Error querying topic embedding: {e}")
            return []

        results = []
        for index in selected_indices:
            results.append(self.topics[index])

        return results


class EmbeddingContext:
    def __init__(self, model):
        self.embeddings = None
        self.content = []
        self.lengths = []
        self.text_data = []
        self.model = model

    def add_resources(self, resources):
        for resource in resources:
            self.content += resource.content
            self.lengths.append(len(self.content))

            embeddings = create_embeddings(resource.content, self.model)

            embeddings = embeddings.reshape(len(resource.content), -1)
            if embeddings.shape[1] != EMBEDDING_DIM:
                print(f"Error embedding resource")
                continue

            if self.embeddings is None:
                self.embeddings = embeddings
            else:
                self.embeddings = torch.cat((self.embeddings, embeddings), dim=0)

            self.text_data.append(resource)

    def query(self, query_text, result_count=1, score_thresh=0.6) -> List[ResearchNote]:
        scores, selected_indices, item_mapping = run_query(
            query_text, self.embeddings, self.model, result_count, score_thresh=score_thresh
        )

        results = []
        for index in selected_indices:
            for i, length in enumerate(self.lengths):
                if index < length:
                    text_data = self.text_data[i]
                    result = ResearchNote(
                        content=self.content[index],
                        title=text_data.title,
                        link=text_data.link,
                        description=text_data.description,
                        outline_items=[k for k in item_mapping.keys() if item_mapping[k] == index]
                    )
                    results.append(result)
                    break

        return results

app/course/embeddings.py
VikParuchuri/textbook_quality
import json
from typing import List

from pydantic import validator
from sqlalchemy import UniqueConstraint
from sqlmodel import JSON, Column, Field, select

from app.components.schemas import AllLessonComponentData
from app.db.base_model import BaseDBModel
from app.db.session import get_session
from app.course.schemas import ResearchNote


class Course(BaseDBModel, table=True):
    __table_args__ = (UniqueConstraint("topic", "model", "version", name="unique_topic_model_version"),)

    model: str
    topic: str
    outline: List[str] = Field(sa_column=Column(JSON), default=list())
    concepts: List[str] = Field(sa_column=Column(JSON), default=list())
    markdown: str = Field(sa_column=Column(JSON), default=list())
    components: List[AllLessonComponentData] = Field(
        sa_column=Column(JSON), default=list()
    )
    queries: List[str] = Field(sa_column=Column(JSON), default=list())
    context: List[ResearchNote] = Field(sa_column=Column(JSON), default=list())
    version: int = Field(default=1, )

    @validator("context")
    def context_to_dict(cls, val: List[ResearchNote]):
        return [v.json() for v in val]

    @validator("components")
    def components_to_dict(cls, val: List[AllLessonComponentData]):
        return [v.json() for v in val]


async def load_cached_course(model: str, topic: str, revision: int):
    async with get_session() as db:
        query = await db.exec(
            select(Course).where(Course.topic == topic, Course.model == model, Course.version == revision)
        )
        course = query.all()
    if len(course) == 0:
        return None
    course = course[0]

    if course.context is not None:
        course.context = [ResearchNote(**json.loads(v)) for v in course.context]

    return course

app/course/models.py
VikParuchuri/textbook_quality
from typing import List

from pydantic import BaseModel


class ResearchNote(BaseModel):
    content: str
    title: str
    link: str
    description: str
    outline_items: List[int]

app/course/schemas.py
VikParuchuri/textbook_quality
from typing import List

from tenacity import RetryError

from app.course.embeddings import EmbeddingContext
from app.course.schemas import ResearchNote
from app.llm.exceptions import GenerationError, InvalidRequestError, RateLimitError
from app.llm.generators.concepts import generate_concepts
from app.llm.generators.outline import generate_outline
from app.services.generators.pdf import download_and_parse_pdfs, search_pdfs
from app.settings import settings
from app.util import debug_print_trace


async def create_course_concepts(course_name: str, revision: int):
    """
    Set the topic and concepts for a course async.
    """
    generated_concepts = None
    try:
        concepts = await generate_concepts(course_name, revision, include_examples=settings.INCLUDE_EXAMPLES)
        if concepts.feasible:
            generated_concepts = concepts.concepts
    except (GenerationError, RateLimitError, InvalidRequestError, RetryError) as e:
        debug_print_trace()
        print(f"Error generating concepts for {course_name}: {e}")

    return generated_concepts


async def create_course_outline(
    course_name: str, concepts: List[str], outline_items: int, revision: int
):
    outline_list = None
    queries = None
    try:
        response = generate_outline(course_name, concepts, revision, item_count=outline_items, include_examples=settings.INCLUDE_EXAMPLES)

        # Stream outline as it generates
        async for outline_data in response:
            outline_list = outline_data.outline
            queries = outline_data.queries
    except (GenerationError, RateLimitError, InvalidRequestError, RetryError) as e:
        debug_print_trace()
        print(f"Error generating outline for {course_name}")

    return outline_list, queries


async def query_course_context(
    model, queries: List[str], outline_items: List[str]
) -> List[ResearchNote] | None:
    # Store the pdf data in the database
    pdf_results = await search_pdfs(queries)
    pdf_data = await download_and_parse_pdfs(pdf_results)

    # If there are no resources, don't generate research notes
    if len(pdf_data) == 0:
        return

    embedding_context = EmbeddingContext(model)
    embedding_context.add_resources(pdf_data)

    results = embedding_context.query(outline_items)
    return results

app/course/tasks.py
VikParuchuri/textbook_quality
from datetime import datetime, timezone
from typing import Optional

import sqlmodel
from sqlalchemy import DateTime, TypeDecorator
from sqlmodel import Field, SQLModel


def get_utc_now():
    return datetime.now(timezone.utc)


class TZDateTime(TypeDecorator):
    impl = DateTime(timezone=True)
    cache_ok = True

    def process_bind_param(self, value, dialect):
        if value is not None:
            value = value.replace(tzinfo=timezone.utc)
        return value

    def process_result_value(self, value, dialect):
        if value is not None:
            value = value.replace(tzinfo=timezone.utc)
        return value


class BaseDBModel(SQLModel):
    id: Optional[int] = Field(primary_key=True, index=True)
    created: datetime | None = Field(
        default_factory=get_utc_now,
        sa_column=sqlmodel.Column(
            TZDateTime,
        ),
    )
    updated: datetime | None = Field(
        default_factory=get_utc_now,
        sa_column=sqlmodel.Column(
            TZDateTime,
            onupdate=get_utc_now,
        ),
    )

app/db/base_model.py
VikParuchuri/textbook_quality
from typing import Generator
from urllib.parse import urlparse, urlunparse

from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import NullPool
from sqlmodel.ext.asyncio.session import AsyncSession

from app.settings import settings


def get_async_db_url():
    conn = settings.DATABASE_URL
    if conn and conn.startswith("postgres://"):
        conn = conn.replace("postgres://", "postgresql+asyncpg://", 1)
    elif conn and conn.startswith("postgresql://"):
        conn = conn.replace("postgresql://", "postgresql+asyncpg://", 1)
    return conn


async_engine = create_async_engine(
    get_async_db_url(),
    future=True,
    pool_pre_ping=True,
    connect_args={"server_settings": {"timezone": "utc"}},
    poolclass=NullPool,
)

async_maker = sessionmaker(
    bind=async_engine,
    class_=AsyncSession,
    expire_on_commit=False,
    autoflush=False,
    autocommit=False,
)


def get_session() -> AsyncSession:
    return async_maker()

app/db/session.py
VikParuchuri/textbook_quality
from app.course.models import Course
from app.db.base_model import BaseDBModel
from app.llm.models import Prompt
from app.services.models import ScrapedData, ServiceResponse

app/db/tables.py
VikParuchuri/textbook_quality
import re
from typing import List

from app.components.schemas import AllLessonComponentData, ComponentNames
from app.lesson.parser import render_lesson_to_markdown


def render_components_to_output_markdown(
    components: List[AllLessonComponentData],
) -> str:
    tuples: List[tuple[str, str]] = []
    for component in components:
        match component.type:
            case ComponentNames.exercise:
                solution = component.solution.strip() if component.solution else ""
                exercise = f"## Exercise\n{component.instructions.strip()}"
                if len(solution) > 0:
                    exercise += f"\n\n### Solution\n{solution}"

                tuples.append(
                    (
                        component.type,
                        exercise,
                    )
                )
            case ComponentNames.section:
                markdown_data = component.markdown.strip()
                if not markdown_data.startswith("#"):
                    markdown_data = f"# {markdown_data}"
                tuples.append((component.type, markdown_data))
            case ComponentNames.text:
                markdown_data = component.markdown.strip()
                tuples.append(
                    (component.type, remove_section_paragraphs(markdown_data))
                )
            case _:
                tuples.append((component.type, component.markdown))
    return render_lesson_to_markdown(tuples, include_type=False)


def remove_section_paragraphs(text):
    # Split text into paragraphs
    paragraphs = text.split("\n")

    # Check if the first paragraph contains "in this section" and remove it if true
    if "in this section" in paragraphs[0].lower():
        paragraphs.pop(0)

    if len(paragraphs) > 0:
        # Check if the last paragraph contains "in the next section" and remove it if true
        if "in the next section" in paragraphs[-1].lower():
            paragraphs.pop()

    # Reconstruct the text from the remaining paragraphs
    replaced = "\n".join(paragraphs).strip()
    replaced = re.sub(r"\n\n+", "\n\n", replaced)
    return replaced

app/lesson/output.py
VikParuchuri/textbook_quality
import re
from typing import List, get_args

from app.components.parser import parse_single_component
from app.components.schemas import COMPONENT_MAP, AllLessonComponentData, ComponentNames
from app.settings import settings


def parse_lesson_markdown(markdown: str) -> List[AllLessonComponentData]:
    # We need to add a section header into the markdown that is generated
    components = []

    lines = markdown.split("\n")
    current_markdown = ""
    current_type = None

    for line in lines:
        if line.startswith("---"):
            if len(current_markdown) > 0 and current_type:
                components.append((current_type, current_markdown.strip()))

            current_type = re.sub("-+", "", line).strip()
            current_markdown = ""
        else:
            current_markdown += line + "\n"
    # Add in the last component
    if len(current_markdown) > 0 and current_type:
        components.append((current_type, current_markdown.strip()))
    return parse_component_list(components)


def parse_component_list(
    components: List[tuple[str, str]]
) -> List[AllLessonComponentData]:
    new_components = []
    for type, markdown in components:
        # Ignore any components that we don't expect
        if type not in list(get_args(settings.VALID_GENERATED_COMPONENTS)):
            continue
        new_components.append(parse_single_component(type, markdown))
    return new_components


def render_lesson_to_markdown(
    components: List[tuple[str, str]], include_type=True
) -> str:
    markdown = ""
    for type, content in components:
        if include_type:
            markdown += f"\n\n---{type}\n\n{content.strip()}"
        else:
            markdown += f"\n\n{content.strip()}"
    return markdown


def render_components_to_markdown(components: List[AllLessonComponentData]) -> str:
    tuples: List[tuple[str, str]] = []
    for component in components:
        match component.type:
            case ComponentNames.exercise:
                # Replace instructions headers if needed
                component.instructions = component.instructions.replace("Instructions\n", "").replace("Instructions:\n", "").strip()
                tuples.append(
                    (
                        component.type,
                        f"Instructions:\n\n{component.instructions}\n\nSolution:\n\n{component.solution}",
                    )
                )
            case _:
                tuples.append((component.type, component.markdown))
    return render_lesson_to_markdown(tuples, include_type=True)

app/lesson/parser.py
VikParuchuri/textbook_quality
from copy import deepcopy
from typing import AsyncGenerator, List

from app.components.parser import generate_component_key
from app.components.schemas import (
    COMPONENT_VERSION,
    AllLessonComponentData,
    ComponentNames,
    MarkdownComponentData,
)
from app.lesson.parser import parse_lesson_markdown, render_components_to_markdown
from app.course.schemas import ResearchNote
from app.llm.exceptions import GenerationError, InvalidRequestError, RateLimitError
from app.llm.generators.lesson import generate_lessons
from app.settings import settings
from app.util import debug_print_trace


async def generate_lesson(
    course_name: str,
    course_components: List[str],
    outline: List[str],
    revision: int,
    research_notes: List[ResearchNote] | None = None,
    sections_per_generation: int = settings.SECTIONS_PER_GENERATION,
) -> List[AllLessonComponentData] | None:
    # Add numbers to the outline - needed for generating the lesson
    numbered_outline = outline

    components = []
    generated_sections = 0
    iterations = 0
    use_cache = True

    while generated_sections < len(numbered_outline) and iterations < len(
        numbered_outline
    ):
        # This is to prime the model with data on what has already been generated
        # The first pass will just include the first section header
        # Subsequent passes will include the previous section and the current section header
        last_section = ""
        if len(components) > 0:
            last_header_index = [
                i for i, c in enumerate(components) if c.type == ComponentNames.section
            ][-1]
            last_section_components = components[last_header_index:]

            # Filter out diagrams, since they take up a lot of tokens.
            # The instructional content should be in the text and examples.
            last_section = render_components_to_markdown(last_section_components)

        current_section_component = MarkdownComponentData(
            type=ComponentNames.section,
            markdown=numbered_outline[generated_sections],
            key=generate_component_key(),
            version=COMPONENT_VERSION,
        )
        components.append(current_section_component)
        current_section_header = render_components_to_markdown(
            [current_section_component]
        )

        current_section = f"{last_section.strip()}\n\n{current_section_header.strip()}"
        current_section = f"{current_section}\n"

        # When to stop generation
        stop_section = None
        if generated_sections + sections_per_generation < len(numbered_outline):
            stop_section = numbered_outline[generated_sections + sections_per_generation]

        # Filter research notes to save tokens, only keep notes relevant to the next 5 sections
        # Find the indices of the next sections
        future_sections = set(list(range(generated_sections, len(numbered_outline)))[:sections_per_generation])
        selected_research_notes = None
        if research_notes is not None:
            selected_research_notes = []
            for research_note in research_notes:
                # If the research note is needed in the next sections
                if set(research_note.outline_items) & future_sections:
                    selected_research_notes.append(research_note)

        try:
            response = generate_single_lesson_chunk(
                numbered_outline,
                current_section,
                generated_sections,
                course_name,
                course_components,
                revision,
                research_notes=selected_research_notes,
                include_examples=settings.INCLUDE_EXAMPLES,
                cache=use_cache,
                stop_section=stop_section,
            )
            new_components = []
            new_component_keys = []
            async for chunk in response:
                new_components = chunk
                # Set keys for the new components to the same as the ones in the last iteration
                for i, key in enumerate(new_component_keys):
                    new_components[i].key = key
                new_component_keys = [c.key for c in new_components]
        except (GenerationError, RateLimitError, InvalidRequestError) as e:
            debug_print_trace()
            print(f"Error generating lesson: {e}")
            return
        except Exception as e:
            debug_print_trace()
            print(f"Error generating lesson: {e}")
            return

        components = deepcopy(components + new_components)

        all_section_headers = [
                i for i, c in enumerate(components) if c.type == ComponentNames.section
            ]
        last_section_index = all_section_headers[-1]

        # Handle partially generated (cut-off) sections
        # Only do this if there are few components, and it's not the end of the lesson
        if len(components) - last_section_index <= 2 and \
                len(components[-1].markdown) < 500 and \
                len(all_section_headers) < len(numbered_outline):
            # If we don't have enough components in the last section, it may have been cut off
            components = components[:last_section_index]
            use_cache = False
        else:
            use_cache = True

        iterations += 1
        generated_sections = len(
            [c for c in components if c.type == ComponentNames.section]
        )
    return components


async def generate_single_lesson_chunk(
    numbered_outline: List[str],
    current_section: str,
    current_section_index: int,
    course_name: str,
    components: List[str],
    revision: int,
    research_notes: List[ResearchNote] | None,
    include_examples: bool,
    cache: bool,
    stop_section: str | None = None,
) -> AsyncGenerator[List[AllLessonComponentData], None]:
    response = generate_lessons(
        numbered_outline,
        current_section,
        current_section_index,
        course_name,
        components,
        revision,
        research_notes=research_notes,
        include_examples=include_examples,
        cache=cache,
        stop_section=stop_section,
    )

    section_start = f"---{ComponentNames.section}"

    async for chunk in response:
        # Remove the final section header from the chunk
        # This happens when we hit the stop token
        if chunk.strip().endswith(section_start):
            chunk = chunk.strip()[:-len(section_start)]
        new_components = parse_lesson_markdown(chunk)
        yield new_components

app/lesson/tasks.py
VikParuchuri/textbook_quality
from copy import deepcopy
from typing import AsyncGenerator, List, Optional

import openai
import stopit
import tiktoken
from aiohttp import ClientPayloadError
from openai.error import APIConnectionError, APIError, ServiceUnavailableError, Timeout

from app.llm.exceptions import GenerationError, InvalidRequestError, RateLimitError
from app.llm.schemas import LLMResponse
from app.settings import settings

openai.api_key = settings.OPENAI_KEY
tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo")
if settings.OPENAI_BASE_URL:
    openai.api_base = settings.OPENAI_BASE_URL


def beginning_of_exception(message: str):
    return message.split(".")[0]


def oai_tokenize_prompt(prompt):
    return len(tokenizer.encode(prompt))


@stopit.threading_timeoutable(default=None)
async def oai_chat_wrapped(
    history: List,
    temperature: float,
    max_tokens: int,
    inner_timeout: int = settings.LLM_TIMEOUT,
    stop_sequences: Optional[List] = None,
    model: str = settings.LLM_TYPE,
) -> AsyncGenerator[str, None]:
    response = await openai.ChatCompletion.acreate(
        model=model,
        messages=history,
        temperature=temperature,
        max_tokens=max_tokens,
        n=1,
        stop=stop_sequences,
        stream=True,
        request_timeout=inner_timeout,
    )
    async for chunk in response:
        stream = chunk
        text = stream["choices"][0]["delta"].get("content", "")
        if text:
            yield text  # Streaming API has the delta property, and the content key inside


@stopit.threading_timeoutable(default=None)
async def oai_prompt_wrapped(
    prompt: str,
    temperature: float,
    max_tokens: int,
    inner_timeout: int = settings.LLM_TIMEOUT,
    stop_sequences: Optional[List] = None,
    model: str = settings.LLM_TYPE,
) -> AsyncGenerator[str, None]:
    response = await openai.Completion.acreate(
        model=model,
        prompt=prompt,
        temperature=temperature,
        max_tokens=max_tokens,
        n=1,
        stop=stop_sequences,
        stream=True,
        request_timeout=inner_timeout,
    )
    async for chunk in response:
        stream = chunk
        text = stream["choices"][0]["text"]
        if text:
            yield text  # Streaming API has the delta property, and the content key inside


async def oai_prompt_response(
    prompt: str,
    temperature: float = settings.LLM_TEMPERATURE,
    timeout: int = settings.LLM_TIMEOUT,
    max_tokens: int = settings.LLM_MAX_RESPONSE_TOKENS,
    stop_sequences=None,
    model: str = settings.LLM_TYPE,
) -> Optional[AsyncGenerator[LLMResponse, None]]:
    response_tokens = 0
    try:
        response = oai_prompt_wrapped(
            prompt,
            temperature,
            max_tokens,
            timeout=timeout,
            inner_timeout=timeout,
            stop_sequences=stop_sequences,
            model=model,
        )
        async for chunk in response:
            response_tokens += 1
            yield LLMResponse(
                text=chunk,
                tokens=response_tokens,
            )
    except (ServiceUnavailableError, APIError, Timeout, APIConnectionError) as e:
        raise GenerationError(beginning_of_exception(str(e)))
    except openai.error.RateLimitError as e:
        raise RateLimitError(beginning_of_exception(str(e)))
    except openai.error.InvalidRequestError as e:
        raise InvalidRequestError(beginning_of_exception(str(e)))
    except ClientPayloadError as e:
        raise GenerationError(beginning_of_exception(str(e)))


async def oai_chat_response(
    prompt: str,
    temperature: float = settings.LLM_TEMPERATURE,
    timeout: int = settings.LLM_TIMEOUT,
    max_tokens: int = settings.LLM_MAX_RESPONSE_TOKENS,
    history=None,
    stop_sequences=None,
    model: str = settings.LLM_TYPE,
) -> Optional[AsyncGenerator[LLMResponse, None]]:
    current_message = {"role": "user", "content": prompt}
    if history is not None:
        history = deepcopy(history)
        history.append(current_message)
    else:
        history = [current_message]

    response_tokens = 0
    try:
        response = oai_chat_wrapped(
            history,
            temperature,
            max_tokens,
            timeout=timeout,
            inner_timeout=timeout,
            stop_sequences=stop_sequences,
            model=model,
        )
        async for chunk in response:
            response_tokens += 1
            yield LLMResponse(
                text=chunk,
                tokens=response_tokens,
            )
    except (ServiceUnavailableError, APIError, Timeout, APIConnectionError) as e:
        raise GenerationError(beginning_of_exception(str(e)))
    except openai.error.RateLimitError as e:
        raise RateLimitError(beginning_of_exception(str(e)))
    except openai.error.InvalidRequestError as e:
        raise InvalidRequestError(beginning_of_exception(str(e)))
    except ClientPayloadError as e:
        raise GenerationError(beginning_of_exception(str(e)))

app/llm/adaptors/oai.py
VikParuchuri/textbook_quality
class GenerationError(Exception):
    pass


class RateLimitError(Exception):
    pass


class InvalidRequestError(Exception):
    pass

app/llm/exceptions.py
VikParuchuri/textbook_quality
import json
import os
from collections import OrderedDict
from json import JSONDecodeError
from typing import List

import ftfy
from pydantic import BaseModel
from tenacity import stop_after_attempt, wait_fixed, before, after, retry, retry_if_exception_type
import threading

from app.llm.exceptions import GenerationError
from app.llm.llm import GenerationSettings, generate_response
from app.llm.prompts import build_prompt
from app.settings import settings
from app.util import extract_only_json_dict


class CourseGeneratedConcepts(BaseModel):
    topic: str
    concepts: List[str]
    feasible: bool


concept_settings = GenerationSettings(
    temperature=0.7,
    max_tokens=256,
    timeout=1200,
    prompt_type="concept",
    model=settings.LLM_INSTRUCT_TYPE,
)


def concept_prompt(topic: str, include_examples=True) -> str:
    with open(os.path.join(settings.EXAMPLE_JSON_DIR, "concepts.json")) as f:
        examples = json.load(f)
    input = OrderedDict([("topic", topic)])
    prompt = build_prompt("concepts", input, examples, include_examples=include_examples)
    return prompt


local_data = threading.local()


def before_retry_callback(retry_state):
    local_data.is_retry = True


def after_retry_callback(retry_state):
    local_data.is_retry = False


@retry(
    retry=retry_if_exception_type(GenerationError),
    stop=stop_after_attempt(5),
    wait=wait_fixed(2),
    before_sleep=before_retry_callback,
    after=after_retry_callback,
    reraise=True,
)
async def generate_concepts(topic: str, revision: int, include_examples: bool = True) -> CourseGeneratedConcepts:
    prompt = concept_prompt(topic, include_examples=include_examples)
    text = ""
    # If we should cache the prompt - skip cache if we're retrying
    should_cache = not getattr(local_data, "is_retry", False)
    response = generate_response(prompt, concept_settings, cache=should_cache, revision=revision)
    async for chunk in response:
        text += chunk
    try:
        text = extract_only_json_dict(text)
        text = str(ftfy.fix_text(text))
        data = json.loads(text.strip())
        concepts = data["concepts"]
        feasible = data["feasible"]
    except (JSONDecodeError, IndexError) as e:
        raise GenerationError(e)

    model = CourseGeneratedConcepts(topic=topic, concepts=concepts, feasible=feasible)
    return model

app/llm/generators/concepts.py
VikParuchuri/textbook_quality
import json
import os
from collections import OrderedDict
from typing import AsyncGenerator, List, get_args

from app.components.schemas import ComponentNames
from app.course.schemas import ResearchNote
from app.llm.llm import GenerationSettings, generate_response
from app.llm.prompts import build_prompt
from app.settings import settings
from copy import deepcopy

lesson_settings = GenerationSettings(
    temperature=0.4,
    max_tokens=6000,
    timeout=1200,
    prompt_type="lesson",
)

COMPONENT_EXTRAS = {
    "exercise": "* ---exercise specifies an exercise. Exercises enable learners to practice the skills learned in the text. They should be challenging and require understanding the text fully (no multiple choice). Exercises should contain instructions and the correct answer. Include 1 exercise block per section.",
    "text": "* ---text specifies a text block. Text blocks are the main content of the course, and teach the learner relevant concepts. Text should be in Github-flavored markdown. Inline math equations should be surrounded with $ symbols, and block math should be surrounded with $$. Include several text blocks per section.",
    "section": "* ---section specifies a section header. Each section matches an item in the outline.",
    "example": "* ---example specifies an example. Examples are used to illustrate concepts found in the text in a practical way.",
}


def lesson_prompt(
    outline: List[str],
    current_section: str,
    current_section_index: int,
    topic: str,
    components: List[str],
    include_examples: bool,
    research_notes: List[ResearchNote] | None = None,
) -> str:
    with open(os.path.join(settings.EXAMPLE_JSON_DIR, "lesson.json")) as f:
        examples = json.load(f)

    # Set default components if none are provided
    if not components:
        components = list(get_args(settings.VALID_GENERATED_COMPONENTS))

    # Strip out components that can't be generated from input
    components = [
        c
        for c in components
        if c in list(get_args(settings.VALID_GENERATED_COMPONENTS))
    ]

    # Basic components that are needed in every lesson
    components += [ComponentNames.text.value, ComponentNames.section.value]
    components = sorted(list(set(components)))

    for example in examples:
        blocks = example["markdown"]
        markdown = ""
        for block_type, content in blocks:
            if block_type in components:
                markdown += f"---{block_type}\n{content}"
        example["markdown"] = markdown

    # Generate a list of extra sentences to be added in to the prompt that are component specific.
    component_extras = [COMPONENT_EXTRAS.get(c, None) for c in components]
    component_extras = [c for c in component_extras if c is not None]

    sections_to_author = len(outline) - current_section_index
    outline_items_to_author = outline[
        current_section_index: current_section_index + sections_to_author
    ]
    outline_items_to_author_str = ",".join(outline_items_to_author)
    outline_stop_item = outline_items_to_author[-1]

    selected_outline = deepcopy(outline)
    if len(outline) > settings.SECTIONS_PER_LESSON:
        surround = min(settings.SECTIONS_PER_LESSON // 2, 10)
        start_item = max(current_section_index - surround, 0)
        end_item = min(current_section_index + sections_to_author + surround, len(outline))
        selected_outline = selected_outline[start_item:end_item]

    rendered_outline = "\n".join(selected_outline).strip()
    items = [
        ("Table of contents\n", rendered_outline),
    ]

    research_notes_exist = research_notes is not None and len(research_notes) > 0
    if research_notes_exist:
        research_content = ""
        for research_note in research_notes:
            content = research_note.content.replace("```", " ")
            content = f"```{content}```"
            research_content += f"* {content}\n"
        items.append(("research notes\n", research_content))

    items.append(("course\n\n", current_section))

    input = OrderedDict(items)

    section_str = "section" if sections_to_author == 1 else "sections"

    prompt = build_prompt(
        "lesson",
        input,
        examples,
        include_examples=include_examples,
        topic=topic,
        components=components,
        component_extras=component_extras,
        research_notes=research_notes_exist,
        section_count=len(outline),
        sections_to_author=sections_to_author,
        section_str=section_str,
        outline_items_to_author=outline_items_to_author_str,
        outline_stop_item=outline_stop_item,
    )
    return prompt


async def generate_lessons(
    outline: List[str],
    current_section: str,
    current_section_index: int,
    topic: str,
    components: List[str],
    revision: int,
    research_notes: List[ResearchNote] | None = None,
    include_examples: bool = True,
    update_after_chars: int = 500,
    cache: bool = True,
    stop_section: str | None = None,
) -> AsyncGenerator[str, None]:
    prompt = lesson_prompt(
        outline,
        current_section,
        current_section_index,
        topic,
        components,
        include_examples,
        research_notes,
    )

    text = ""
    stop_sequences = None
    if stop_section is not None:
        stop_sequences = [stop_section]

    response = generate_response(prompt, lesson_settings, cache=cache, revision=revision, stop_sequences=stop_sequences)
    chunk_len = 0

    # Yield text in batches, to avoid creating too many DB models
    async for chunk in response:
        text += chunk
        chunk_len += len(chunk)
        if chunk_len >= update_after_chars:
            yield text
            chunk_len = 0
    # Yield the remaining text
    yield text

app/llm/generators/lesson.py
VikParuchuri/textbook_quality
import json
import os
import re
import threading
from collections import OrderedDict
from json import JSONDecodeError
from typing import AsyncGenerator, List

import ftfy
from pydantic import BaseModel, parse_obj_as
from tenacity import retry_if_exception_type, stop_after_attempt, retry, wait_fixed

from app.llm.exceptions import GenerationError
from app.llm.llm import GenerationSettings, generate_response
from app.llm.prompts import build_prompt
from app.settings import settings
from app.util import extract_only_json_dict

outline_settings = GenerationSettings(
    temperature=0.6,
    max_tokens=2048,
    timeout=1200,
    prompt_type="outline",
    model=settings.LLM_INSTRUCT_TYPE,
)

# This can get better results from a finetuned model, forces a certain outline format
prompt_start_hint = '\n{"outline": ["1. '

class GeneratedOutlineData(BaseModel):
    outline: List[str]
    queries: List[str] | None = None


def outline_prompt(topic: str, concepts: List[str], item_count: int = settings.SECTIONS_PER_LESSON, include_examples=True) -> str:
    with open(os.path.join(settings.EXAMPLE_JSON_DIR, "outline.json")) as f:
        examples = json.load(f)
    input = OrderedDict([("topic", topic), ("concepts", concepts)])
    prompt = build_prompt(
        "outline",
        input,
        examples,
        topic=topic,
        concepts=concepts,
        item_count=item_count,
        include_examples=include_examples,
    )
    if settings.FINETUNED:
        prompt += prompt_start_hint
    return prompt


def parse_json_data(outline: dict) -> GeneratedOutlineData:
    outline = parse_obj_as(GeneratedOutlineData, outline)
    # Get rid of prefix numbers if they exist (they are sometimes added, but we want to strip them out for consistency)
    return outline


def try_parse_json(text: str) -> dict | None:
    data = None
    try:
        data = json.loads(text.strip())
    except json.decoder.JSONDecodeError:
        # Try to re-parse if it failed
        try:
            data = json.loads(text.strip() + '"]}')
        except json.decoder.JSONDecodeError:
            # If it fails again, keep going with the loop.
            pass
    return data


local_data = threading.local()


def before_retry_callback(retry_state):
    local_data.is_retry = True


def after_retry_callback(retry_state):
    local_data.is_retry = False


@retry(
    retry=retry_if_exception_type(GenerationError),
    stop=stop_after_attempt(5),
    wait=wait_fixed(2),
    before_sleep=before_retry_callback,
    after=after_retry_callback,
    reraise=True,
)
async def generate_outline(
    topic: str,
    concepts: List[str],
    revision: int,
    update_after_chars: int = 50,
    item_count: int = 10,
    include_examples: bool = True
) -> AsyncGenerator[GeneratedOutlineData, None]:
    # Sort concepts alphabetically so that the prompt is the same every time
    concepts = sorted(concepts)
    prompt = outline_prompt(topic, concepts, item_count=item_count, include_examples=include_examples)
    text = ""
    if settings.FINETUNED:
        text = prompt_start_hint
    # Do not hit cache on retries
    should_cache = not getattr(local_data, "is_retry", False)
    response = generate_response(prompt, outline_settings, cache=should_cache, revision=revision)

    chunk_len = 0
    async for chunk in response:
        text += chunk
        chunk_len += len(chunk)
        if chunk_len >= update_after_chars:
            data = try_parse_json(text.strip())
            if data:
                yield parse_json_data(data)
            chunk_len = 0

    # Handle the last bit of data
    try:
        # Strip out text before/after the json.  Sometimes the LLM will include something before the json input.
        text = extract_only_json_dict(text)
        text = str(ftfy.fix_text(text))
        data = json.loads(text.strip())
    except JSONDecodeError as e:
        raise GenerationError(e)
    yield parse_json_data(data)


def renumber_outline(outline):
    def renumber(match):
        major, minor = int(match.group(1)), match.group(2)
        major -= 1
        return f"{major}{minor}"

    new_outline = [re.sub(r"(\d+)(\..*|$)", renumber, chapter) for chapter in outline]

    return new_outline
app/llm/generators/outline.py
VikParuchuri/textbook_quality
import json
import os
from collections import OrderedDict
from json import JSONDecodeError
from typing import List

from app.llm.exceptions import GenerationError
from app.llm.llm import GenerationSettings, generate_response
from app.llm.prompts import build_prompt
from app.settings import settings
from app.util import extract_only_json_list

title_settings = GenerationSettings(
    temperature=0.9,
    max_tokens=512,
    timeout=40,
    prompt_type="title"
)


def title_prompt(subject: str) -> str:
    with open(os.path.join(settings.EXAMPLE_JSON_DIR, f"title.json")) as f:
        examples = json.load(f)
    input = OrderedDict([("subject", subject)])
    prompt = build_prompt("title", input, examples)
    return prompt


async def generate_title(
    subject: str,
) -> List[str]:
    prompt = title_prompt(subject)
    text = ""
    response = generate_response(prompt, title_settings, cache=False)
    async for chunk in response:
        text += chunk

    try:
        text = extract_only_json_list(text)
        data = json.loads(text.strip())
    except (JSONDecodeError, IndexError) as e:
        raise GenerationError(e)
    return data

app/llm/generators/title.py
VikParuchuri/textbook_quality
import json
import os
from collections import OrderedDict
from copy import deepcopy
from json import JSONDecodeError
from typing import List

import ftfy
from pydantic import BaseModel

from app.llm.exceptions import GenerationError
from app.llm.llm import GenerationSettings, generate_response
from app.llm.prompts import build_prompt
from app.settings import settings
from app.util import extract_only_json_dict
from app.llm.adaptors.oai import oai_tokenize_prompt


class GeneratedTOC(BaseModel):
    topic: str
    outline: List[str]
    queries: List[str]


toc_settings = GenerationSettings(
    temperature=0.5,
    max_tokens=1024,
    timeout=1200,
    prompt_type="toc",
    model=settings.LLM_TYPE,
)


def toc_prompt(topic: str, toc: str, include_examples=True) -> str:
    with open(os.path.join(settings.EXAMPLE_JSON_DIR, "toc.json")) as f:
        examples = json.load(f)
    input = OrderedDict([
        ("topic", topic),
        ("toc", toc),
    ])
    prompt = build_prompt("toc", input, examples, include_examples=include_examples)
    return prompt


async def generate_tocs(topic: str, draft_toc: str, include_examples: bool = True) -> GeneratedTOC | None:
    prompt = toc_prompt(topic, draft_toc, include_examples=include_examples)
    text = ""

    settings_inst = deepcopy(toc_settings)
    try:
        settings_inst.max_tokens = oai_tokenize_prompt(draft_toc) + 512 # Max tokens to generate
    except Exception:
        return

    response = generate_response(prompt, settings_inst)
    async for chunk in response:
        text += chunk
    try:
        text = extract_only_json_dict(text)
        text = str(ftfy.fix_text(text))
        data = json.loads(text.strip())
        toc = data["outline"]
        queries = data["queries"]
    except (JSONDecodeError, IndexError) as e:
        raise GenerationError(e)

    model = GeneratedTOC(topic=topic, outline=toc, queries=queries)
    return model

app/llm/generators/toc.py
VikParuchuri/textbook_quality
import json
import os
from collections import OrderedDict
from json import JSONDecodeError
from typing import List, Optional

from app.llm.exceptions import GenerationError
from app.llm.llm import GenerationSettings, generate_response
from app.llm.prompts import build_prompt
from app.settings import settings
from app.util import extract_only_json_list

topic_settings = GenerationSettings(
    temperature=0.9,
    max_tokens=512,
    timeout=40,
    prompt_type="topic",
    model=settings.LLM_INSTRUCT_TYPE,
)


def topic_prompt(book_title: str) -> str:
    with open(os.path.join(settings.EXAMPLE_JSON_DIR, f"topic.json")) as f:
        examples = json.load(f)
    input = OrderedDict([("title", book_title)])
    prompt = build_prompt("topic", input, examples, title=book_title)
    return prompt


async def generate_topic(
    book_title: str,
) -> List[str]:
    prompt = topic_prompt(book_title)
    text = ""
    response = generate_response(prompt, topic_settings)
    async for chunk in response:
        text += chunk
    try:
        text = extract_only_json_list(text)
        data = json.loads(text.strip())
    except JSONDecodeError as e:
        raise GenerationError(e)
    return data


def topic_specific_prompt(book_title: str, domain: Optional[str]) -> str:
    with open(os.path.join(settings.EXAMPLE_JSON_DIR, f"specific_topic.json")) as f:
        examples = json.load(f)
    input = OrderedDict([("title", book_title)])
    prompt = build_prompt(
        "specific_topic", input, examples, title=book_title, domain=domain
    )
    return prompt


async def generate_specific_topic(
    book_title: str,
    domain: Optional[str] = None,
) -> List[str]:
    prompt = topic_specific_prompt(book_title, domain)
    text = ""
    response = generate_response(prompt, topic_settings)
    async for chunk in response:
        text += chunk

    try:
        text = extract_only_json_list(text)
        data = json.loads(text.strip())
    except JSONDecodeError as e:
        raise GenerationError(e)
    return data

app/llm/generators/topic.py
VikParuchuri/textbook_quality
import asyncio
import hashlib
import time
from typing import AsyncGenerator, List, Optional

from sqlalchemy.exc import IntegrityError
from sqlmodel import select

from app.db.session import get_session
from app.llm.adaptors.oai import (
    oai_chat_response,
    oai_prompt_response,
    oai_tokenize_prompt,
)
from app.llm.exceptions import GenerationError, InvalidRequestError, RateLimitError
from app.llm.models import Prompt
from app.llm.schemas import GenerationSettings
from app.settings import settings
from app.util import fix_unicode_text


async def generate_response(
    prompt: str,
    prompt_settings: GenerationSettings,
    history: Optional[List] = None,
    max_tries: int = 2,
    cache: bool = True,
    revision: int = 1,
    stop_sequences: Optional[List[str]] = None,
) -> AsyncGenerator[str, None]:
    temperature = prompt_settings.temperature
    max_tokens = prompt_settings.max_tokens
    timeout = prompt_settings.timeout
    prompt_stops = prompt_settings.stop_sequences
    prompt_type = prompt_settings.prompt_type
    model = (
        prompt_settings.model or settings.LLM_TYPE
    )  # Use default model if not specified

    # Stop sequences for the llm
    stops = []
    if prompt_stops is not None:
        stops.extend(prompt_stops)
    if stop_sequences is not None:
        stops.extend(stop_sequences)

    # Only support up to 4 stop sequences
    if len(stops) == 0:
        stops = None
    else:
        stops = stops[:4]

    # Remove utf-8 surrogate characters
    prompt = fix_unicode_text(prompt)

    # Hash the prompt as a DB key
    hash = hashlib.sha512()
    hash.update(prompt.encode("utf-8"))
    hex = hash.hexdigest()

    async with get_session() as db:
        # Break if we've already run this prompt
        query = await db.exec(
            select(Prompt).where(Prompt.hash == hex, Prompt.model == settings.LLM_TYPE, Prompt.version == revision)
        )
        prompt_model = query.first()


    if prompt_model is not None and cache:
        yield prompt_model.response
        return

    orig_model = model
    for i in range(max_tries):
        try:
            match model:
                case "gpt-3.5-turbo" | "gpt-4":
                    prompt_tokens = oai_tokenize_prompt(prompt)

                    # Reduce tokens requested if we have too many in the prompt
                    if (
                        prompt_tokens + max_tokens
                        >= settings.LLM_TYPES[model]["max_tokens"]
                    ):
                        # Use extended model if we have too many tokens
                        model = settings.LLM_EXTENDED_TYPE
                        if (
                            prompt_tokens + max_tokens
                            >= settings.LLM_TYPES[model]["max_tokens"]
                        ):
                            raise InvalidRequestError(
                                f"Input prompt is too long, requested {prompt_tokens} prompt tokens and {max_tokens} generation tokens."
                            )

                    response = oai_chat_response(
                        prompt,
                        temperature,
                        timeout,
                        max_tokens,
                        history,
                        stops,
                        model=model,
                    )
                case "gpt-3.5-turbo-instruct":
                    prompt_tokens = oai_tokenize_prompt(prompt)
                    if (
                        prompt_tokens + max_tokens
                        >= settings.LLM_TYPES[model]["max_tokens"]
                    ):
                        raise InvalidRequestError(
                            f"Input prompt is too long, requested {prompt_tokens} prompt tokens and {max_tokens} generation tokens."
                        )

                    response = oai_prompt_response(
                        prompt,
                        temperature,
                        timeout,
                        max_tokens,
                        stops,
                        model=model,
                    )
                case _:
                    if model not in settings.LLM_TYPES:
                        raise NotImplementedError(
                            "This LLM type is not supported currently."
                        )

                    prompt_tokens = oai_tokenize_prompt(prompt)

                    allowed_tokens = settings.LLM_TYPES[model]["max_tokens"]
                    if prompt_tokens + max_tokens > allowed_tokens:
                        max_tokens = allowed_tokens - prompt_tokens

                    if max_tokens < 256:
                        raise InvalidRequestError(
                            f"Input prompt is too long, requested {prompt_tokens} prompt tokens and {max_tokens} generation tokens."
                        )

                    response = oai_prompt_response(
                        prompt,
                        temperature,
                        timeout,
                        max_tokens,
                        stops,
                        model=model,
                    )
            break
        except (GenerationError, RateLimitError, InvalidRequestError):
            # Re-raise error if we're on the last try
            model = orig_model
            if i == max_tries - 1:
                raise

            await asyncio.sleep(30 * (i + 1))

    full_text = ""
    async for chunk in response:
        text = chunk.text
        yield text
        full_text += text

    # Skip caching
    if not cache:
        return

    async with get_session() as db:
        try:
            prompt_model = Prompt(
                hash=hex,
                prompt=prompt,
                response=full_text,
                type=prompt_type,
                model=settings.LLM_TYPE,
                version=revision,
            )
            db.add(prompt_model)
            await db.commit()
        except IntegrityError:
            await db.rollback()

app/llm/llm.py
VikParuchuri/textbook_quality
from sqlmodel import Field, UniqueConstraint

from app.db.base_model import BaseDBModel
from app.util import BaseEnum


class PromptTypes(str, BaseEnum):
    lesson = "lesson"
    concepts = "concept"
    outline = "outline"
    topic = "topic"
    title = "title"
    toc = "toc"


class Prompt(BaseDBModel, table=True):
    __table_args__ = (UniqueConstraint("hash", "model", "version", name="unique_hash_model_version"),)
    hash: str = Field(index=True)
    prompt: str
    response: str
    type: PromptTypes
    model: str
    version: int = Field(default=1)

app/llm/models.py
VikParuchuri/textbook_quality
import json
from collections import OrderedDict
from typing import List, Optional

from jinja2 import Environment, FileSystemLoader, Template

from app.settings import settings


def render_single_dict(d: OrderedDict) -> str:
    prompt = ""
    for k in d.keys():
        if k in ["json", "markdown"]:
            continue
        header = k.replace("_", " ").title()
        if not header.endswith((":", "?", ".", "#", "\n")):
            header += ":"
        if not header.endswith(("\n", " ")):
            header += " "
        prompt += f"{header}{str(d[k])}\n"
    if "json" in d:
        prompt += json.dumps(d["json"])
    if "markdown" in d:
        prompt += f"\n{d['markdown']}"
    return prompt


def render_examples(examples: List[OrderedDict]) -> str:
    if examples is None:
        return ""
    example_prompt = ""
    for i, example in enumerate(examples):
        example_prompt += f"Example {i + 1}\n"
        example_prompt += render_single_dict(example)
        example_prompt += "\n"
    return example_prompt


def render_input(input: OrderedDict) -> str:
    if input is None:
        return ""
    input_prompt = "Input\n"
    input_prompt += render_single_dict(input)
    return input_prompt


def build_prompt(
    template_name: str,
    input: OrderedDict,
    examples: Optional[list] = None,
    include_examples: bool = True,
    **keys,
) -> str:
    with open(f"{settings.PROMPT_TEMPLATE_DIR}/{template_name}.jinja") as file_:
        template_str = file_.read()
    template = Environment(
        loader=FileSystemLoader(settings.PROMPT_TEMPLATE_DIR)
    ).from_string(template_str)
    instruction = template.render(**keys)
    input_prompt = render_input(input)
    if include_examples:
        example_prompt = render_examples(examples)
        prompt = f"{instruction}\n\n{example_prompt}\n{input_prompt}"
    else:
        prompt = f"{instruction}\n\n{input_prompt}"
    return prompt

app/llm/prompts.py
VikParuchuri/textbook_quality
from typing import List, Optional

from pydantic import BaseModel


class LLMResponse(BaseModel):
    text: str
    tokens: int


class GenerationSettings(BaseModel):
    temperature: float
    max_tokens: int
    timeout: int
    stop_sequences: Optional[List[str]]
    prompt_type: str
    component_name: Optional[str]
    model: Optional[str]

app/llm/schemas.py
VikParuchuri/textbook_quality
import urllib.parse
from json import JSONDecodeError

import aiohttp

from app.services.exceptions import RequestError, ResponseError
from app.services.schemas import ServiceInfo, ServiceNames, ServiceSettings
from app.settings import settings

serpapi_pdf_search_settings = ServiceSettings(name=ServiceNames.serpapi, type="pdf")


async def serpapi_router(service_settings: ServiceSettings, service_info: ServiceInfo):
    match service_settings.type:
        case "google":
            response = await run_search(
                service_info.query, "search.json", query_params="&engine=google"
            )
        case "pdf":
            query = f"{service_info.query} filetype:pdf"
            response = await run_search(
                query, "search.json", query_params="&engine=google&gl=us&hl=en"
            )
        case _:
            raise RequestError(f"Unknown serpapi service type {service_settings.type}")
    return response


async def run_search(query: str, endpoint: str, query_params: str):
    encoded = urllib.parse.quote_plus(query)

    request_url = f"https://serpapi.com/{endpoint}?q={encoded}{query_params}&safe=active&api_key={settings.SERPAPI_KEY}"

    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(request_url) as response:
                json = await response.json()
    except aiohttp.ClientResponseError as e:
        raise RequestError(f"Serpapi request failed with status {e.status}")
    except JSONDecodeError as e:
        raise ResponseError(f"Could not decode Serpapi response as JSON: {e}")

    return json

app/services/adaptors/serpapi.py
VikParuchuri/textbook_quality
import urllib.parse
from json import JSONDecodeError

import aiohttp

from app.services.exceptions import RequestError, ResponseError
from app.services.schemas import ServiceInfo, ServiceNames, ServiceSettings
from app.settings import settings

serply_pdf_search_settings = ServiceSettings(name=ServiceNames.serply, type="pdf")


async def serply_router(service_settings: ServiceSettings, service_info: ServiceInfo):
    match service_settings.type:
        case "google":
            response = await run_search(service_info.query, "search")
        case "pdf":
            response = await run_search(
                service_info.query, "search", "&as_filetype=pdf"
            )
        case _:
            raise RequestError(f"Unknown Serply service type {service_settings.type}")
    return response


async def run_search(query: str, endpoint: str, query_params: str = ""):
    encoded = urllib.parse.quote_plus(query)

    request_url = f"https://api.serply.io/v1/{endpoint}/q={encoded}&num=10&safe=active{query_params}"
    headers = {"X-Api-Key": settings.SERPLY_KEY, "X-Proxy-Location": "US"}

    try:
        async with aiohttp.ClientSession(headers=headers) as session:
            async with session.get(request_url) as response:
                json = await response.json()
    except (aiohttp.ClientResponseError, aiohttp.ClientOSError) as e:
        raise RequestError(f"Request failed with status {e.status}")
    except JSONDecodeError as e:
        raise ResponseError(f"Could not decode response as JSON: {e}")

    if json["total"] == 0:
        raise ResponseError("Nothing found")

    return json

app/services/adaptors/serply.py
VikParuchuri/textbook_quality
from typing import List, Optional

from sqlmodel import select
from sqlmodel.ext.asyncio.session import AsyncSession

from app.db.session import get_session
from app.services.models import ScrapedData, ServiceResponse


async def get_stored_urls(urls: List[str]) -> List[Optional[str]]:
    async with get_session() as db:
        query = await db.exec(select(ScrapedData).where(ScrapedData.source.in_(urls)))
        stored_urls = query.all()
    return_data = {}
    for url in urls:
        return_data[url] = None
        for stored_url in stored_urls:
            if url == stored_url.source:
                return_data[url] = stored_url.uploaded
    return [return_data[url] for url in urls]


async def get_service_response_model(name: str, hex: str):
    async with get_session() as db:
        query = await db.exec(
            select(ServiceResponse).where(
                ServiceResponse.hash == hex, ServiceResponse.name == name
            )
        )
        service_model = query.first()
    return service_model

app/services/dependencies.py
VikParuchuri/textbook_quality
class RequestError(Exception):
    pass


class ResponseError(Exception):
    pass


class ProcessingError(Exception):
    pass

app/services/exceptions.py
VikParuchuri/textbook_quality
import asyncio
import math
import os
from itertools import chain
from typing import List, Optional

import fitz as pymupdf
from aiohttp import ClientPayloadError
from fitz import FileDataError
from pydantic import BaseModel

from app.db.session import get_session
from app.services.adaptors.serpapi import serpapi_pdf_search_settings
from app.services.adaptors.serply import serply_pdf_search_settings
from app.services.dependencies import get_stored_urls
from app.services.exceptions import ProcessingError
from app.services.models import store_scraped_data
from app.services.network import download_and_save
from app.services.schemas import PDFData, ServiceInfo, ServiceNames
from app.services.service import get_service_response
from app.settings import settings

BLOCK_SIZE = 2000  # Characters per text block

SEARCH_SETTINGS = {
    "serply": serply_pdf_search_settings,
    "serpapi": serpapi_pdf_search_settings,
}

pdf_service_settings = SEARCH_SETTINGS.get(
    settings.SEARCH_BACKEND
)  # Get the right settings for search

if not pdf_service_settings:
    print("Not using retrieval.  Invalid search backend or no backend set.")


class PDFSearchResult(BaseModel):
    link: str
    title: str
    description: str
    query: str


async def search_pdfs(queries: List[str], max_count=3) -> List[PDFSearchResult]:
    coroutines = [search_pdf(query, max_count) for query in queries]

    # Run queries sequentially
    results = []
    for routine in coroutines:
        results.append(await routine)

    results = list(chain.from_iterable(results))

    # Filter results to only unique links
    filtered = []
    seen_links = []
    seen_titles = []
    for r in results:
        link = r.link
        title = r.title
        if link not in seen_links and title not in seen_titles:
            seen_links.append(link)
            seen_titles.append(title)
            filtered.append(r)
    return filtered


async def search_pdf(query: str, max_count) -> List[PDFSearchResult]:
    if not pdf_service_settings:
        return []

    service_info = ServiceInfo(query=query)

    # Run google search for pdf
    response = await get_service_response(pdf_service_settings, service_info)

    pdf_links = []
    for i in range(10):
        try:
            if pdf_service_settings.name == ServiceNames.serply:
                pdf_result = response["results"][i]
                pdf_description = pdf_result["description"]
            else:
                pdf_result = response["organic_results"][i]
                pdf_description = pdf_result["snippet"]

            pdf_link = pdf_result["link"]
            pdf_title = pdf_result["title"]
            if not pdf_link or not pdf_link.endswith(".pdf"):
                continue
        except (KeyError, IndexError):
            continue

        search_result = PDFSearchResult(
            link=pdf_link, title=pdf_title, description=pdf_description, query=query
        )
        pdf_links.append(search_result)
        if len(pdf_links) >= max_count:
            break
    return pdf_links


async def download_and_parse_pdfs(
    search_results: List[PDFSearchResult],
) -> List[PDFData]:
    # Deduplicate links
    deduped_search_results = []
    seen_links = set()
    for search_result in search_results:
        if search_result.link not in seen_links:
            deduped_search_results.append(search_result)
            seen_links.add(search_result.link)

    links = [search_result.link for search_result in deduped_search_results]
    pdf_paths = await get_stored_urls(links)

    coroutines = [
        download_and_parse_pdf(search_result, pdf_path) for search_result, pdf_path in zip(deduped_search_results, pdf_paths)
    ]
    results = await asyncio.gather(*coroutines)
    results = [result for result in results if result]

    # Store all scraping results
    async with get_session() as db:
        for result in results:
            if not result.stored:
                store_scraped_data(db, result.link, result.pdf_path)
        await db.commit()
    return results


async def download_and_parse_pdf(search_result: PDFSearchResult, pdf_path: Optional[str]) -> Optional[PDFData]:
    stored = False
    if pdf_path:
        with open(os.path.join(settings.PDF_CACHE_DIR, pdf_path), "rb") as f:
            pdf_data = f.read()
        stored = True
    else:
        try:
            # Download pdf, save to filesystem
            pdf_path, pdf_data = await download_and_save(search_result.link)
        except ProcessingError as e:
            # This happens when pdf is too large
            return
        except ClientPayloadError as e:
            return
        except Exception as e:
            return

    try:
        pdf_content = parse_pdf(pdf_data)
    except FileDataError:
        return

    pdf_cls = PDFData(
        pdf_path=pdf_path,
        link=search_result.link,
        title=search_result.title,
        description=search_result.description,
        content=pdf_content,
        query=search_result.query,
        stored=stored,
    )

    # Bail out if we don't have enough text or blocks in our pdf
    if len(pdf_cls.content) < 6:
        return

    all_content = "".join(pdf_cls.content)
    if len(all_content) < 2000:
        return

    return pdf_cls


def parse_pdf(data) -> List[str]:
    with pymupdf.open(stream=data) as doc:
        blocks = []
        for page in doc:
            blocks += page.get_text(
                "blocks",
                sort=True,
                flags=~pymupdf.TEXT_PRESERVE_LIGATURES
                & pymupdf.TEXT_PRESERVE_WHITESPACE
                & ~pymupdf.TEXT_PRESERVE_IMAGES
                & ~pymupdf.TEXT_INHIBIT_SPACES
                & pymupdf.TEXT_DEHYPHENATE
                & pymupdf.TEXT_MEDIABOX_CLIP,
            )
    start = math.floor(len(blocks) * 0.15)
    end = math.ceil(len(blocks) * 0.85)

    blocks = blocks[start:end]
    parsed_blocks = []
    block = ""
    for i, b in enumerate(blocks):
        block += b[4]
        if len(block) > BLOCK_SIZE:
            parsed_blocks.append(block)
            block = ""
    parsed_blocks.append(block)
    return parsed_blocks

app/services/generators/pdf.py
VikParuchuri/textbook_quality
from pydantic import validator
from sqlmodel import JSON, Column, Field, UniqueConstraint, Session
from sqlmodel.ext.asyncio.session import AsyncSession

from app.db.base_model import BaseDBModel
from app.db.session import get_session
from app.services.schemas import ServiceInfo, ServiceNames


class ServiceResponse(BaseDBModel, table=True):
    __table_args__ = (UniqueConstraint("hash", "name", name="unique_hash_name"),)
    hash: str = Field(index=True)
    request: ServiceInfo = Field(sa_column=Column(JSON), default=dict(), nullable=False)
    response: dict = Field(sa_column=Column(JSON), default=dict(), nullable=False)
    name: ServiceNames

    @validator("request")
    def gen_request_to_dict(cls, val: ServiceInfo):
        return val.json()


class ScrapedData(BaseDBModel, table=True):
    source: str = Field(index=True, unique=True)
    uploaded: str
    extra: dict | None = Field(sa_column=Column(JSON), default=dict(), nullable=True)


def store_scraped_data(db: AsyncSession, source: str, uploaded: str):
    data = ScrapedData(source=source, uploaded=uploaded)
    db.add(data)

app/services/models.py
VikParuchuri/textbook_quality
import os
import secrets
from io import BytesIO

import aiohttp
from aiohttp.client_exceptions import ClientConnectorError, ClientOSError

from app.services.exceptions import ProcessingError
from app.settings import settings


def generate_pdf_name():
    return f"{secrets.token_urlsafe(32)}.pdf"


async def download_file_safely(url: str) -> bytes:
    try:
        data = BytesIO()
        total_len = 0
        chunk_size = 1024
        async with aiohttp.ClientSession() as session:
            async with session.get(url) as response:
                content_length = int(response.headers.get("content-length", 0))
                if content_length > settings.MAX_DOWNLOAD_SIZE:
                    raise ProcessingError(f"File too large: {url}")

                async for chunk in response.content.iter_chunked(chunk_size):
                    total_len += chunk_size
                    data.write(chunk)
                    if total_len > settings.MAX_DOWNLOAD_SIZE:
                        raise ProcessingError(f"Download exceeded max size: {url}")
    except (ConnectionError, ClientConnectorError, ClientOSError):
        raise ProcessingError(f"Failed to download file: {url}")

    return data.getvalue()


async def download_and_save(url: str):
    data = await download_file_safely(url)

    pdf_name = generate_pdf_name()
    file_path = os.path.join(settings.PDF_CACHE_DIR, pdf_name)
    with open(file_path, "wb") as f:
        f.write(data)
    return pdf_name, data

app/services/network.py
VikParuchuri/textbook_quality
from typing import List, Optional

from pydantic import BaseModel

from app.util import BaseEnum


class ServiceNames(str, BaseEnum):
    serply = "serply"
    serpapi = "serpapi"


class ServiceSettings(BaseModel):
    name: ServiceNames
    type: str


class ServiceInfo(BaseModel):
    query: Optional[str] = None
    content: Optional[str] = None


class PDFData(BaseModel):
    pdf_path: str
    link: str
    text_link: str | None = None
    title: str
    description: str
    content: List[str]
    query: str | None = None
    stored: bool = False

app/services/schemas.py
VikParuchuri/textbook_quality
import hashlib

from app.db.session import get_session
from app.services.adaptors.serpapi import serpapi_router
from app.services.adaptors.serply import serply_router
from app.services.dependencies import get_service_response_model
from app.services.models import ServiceResponse
from app.services.schemas import ServiceInfo, ServiceNames, ServiceSettings


async def get_service_response(
    service_settings: ServiceSettings,
    service_info: ServiceInfo,
) -> dict:
    hash = hashlib.sha512()
    # Turn dict into list, sort keys, then hash.  This ensures consistent order.
    service_info_str = str(sorted(service_info.dict().items())).encode("utf-8")
    hash.update(service_info_str)
    hex = hash.hexdigest()

    # Break if we've already run this query
    service_model = await get_service_response_model(service_settings.name, hex)

    if service_model is not None:
        return service_model.response

    match service_settings.name:
        case ServiceNames.serply:
            response = await serply_router(service_settings, service_info)
        case ServiceNames.serpapi:
            response = await serpapi_router(service_settings, service_info)
        case _:
            raise NotImplementedError("This Service type is not currently supported.")

    async with get_session() as db:
        # Save the response to the DB
        service_model = ServiceResponse(
            hash=hex,
            request=service_info.dict(),
            response=response,
            name=service_settings.name,
        )
        db.add(service_model)
        await db.commit()

    return response

app/services/service.py
VikParuchuri/textbook_quality
import os
from typing import Literal, Optional

from dotenv import find_dotenv
from pydantic import BaseSettings


class Settings(BaseSettings):
    # Path settings
    BASE_DIR: str = os.path.abspath(os.path.dirname(__file__))
    PDF_CACHE_DIR = os.path.join(BASE_DIR, "cache")  # Where to save pdfs
    DATA_DIR = os.path.join(BASE_DIR, "data")  # Where to save data
    PROMPT_TEMPLATE_DIR: str = os.path.join(BASE_DIR, "llm", "templates")
    EXAMPLE_JSON_DIR: str = os.path.join(BASE_DIR, "llm", "examples")

    # Database
    DATABASE_URL: str = "postgresql://localhost/textbook"
    DEBUG: bool = False

    # Content
    SECTIONS_PER_LESSON: int = 30  # Lower this to make books shorter
    SECTIONS_PER_GENERATION: int = 5 # How many sections to generate in one prompt
    MAX_DOWNLOAD_SIZE: int = 6 * 1024 * 1024  # Max pdf size to download, 6 MB
    FINETUNED: bool = False # If we're using a finetuned textbook gen model
    INCLUDE_EXAMPLES: bool = (
        True  # Include examples in prompts, False with custom model
    )

    # LLM
    LLM_TYPES = {
        "gpt-3.5-turbo": {"max_tokens": 4097},
        "gpt-3.5-turbo-16k": {"max_tokens": 16384},
        "gpt-3.5-turbo-instruct": {"max_tokens": 4097},
        "llama": {"max_tokens": 8192},
        "gpt-4": {"max_tokens": 8192},
        "gpt-4-32k": {"max_tokens": 32768},
    }

    LLM_TEMPERATURE: float = 0.5
    LLM_TIMEOUT: int = 480
    LLM_MAX_RESPONSE_TOKENS: int = 2048
    OPENAI_KEY: str = ""
    OPENAI_BASE_URL: Optional[str] = None
    LLM_TYPE: str = "gpt-3.5-turbo"
    LLM_INSTRUCT_TYPE: str = "gpt-3.5-turbo-instruct"
    LLM_EXTENDED_TYPE: str = "gpt-3.5-turbo-16k"

    # Generation
    VALID_GENERATED_COMPONENTS = Literal["text", "example", "exercise", "section"]
    VALID_COMPONENTS = Literal["text", "example", "exercise", "section"]

    # Retrieval backend service
    SERPLY_KEY: str = ""
    SERPAPI_KEY: str = ""
    SEARCH_BACKEND: Optional[str] = "serply"

    # General
    THREADS_PER_WORKER: int = 1 # How many threads to use per worker process to save RAM
    RAY_CACHE_PATH: Optional[str] = None # Where to save ray cache
    RAY_DASHBOARD_HOST: str = "0.0.0.0"
    RAY_CORES_PER_WORKER = .5 # How many cpu cores to allocate per worker

    class Config:
        env_file = find_dotenv("local.env")


settings = Settings()

app/settings.py
VikParuchuri/textbook_quality
import traceback
from enum import Enum
import regex

import ftfy

from app.settings import settings


class BaseEnum(Enum):
    def __str__(self):
        return str(self.value)


def extract_only_json_dict(text: str) -> str:
    # Extract the first top-level JSON object
    pattern = r'({(?:[^{}]|(?R))*})'
    match = regex.search(pattern, text, regex.DOTALL)
    if match:
        return match.group(0)

    return text


def extract_only_json_list(text: str) -> str:
    # Extract the first top-level JSON object
    pattern = r'(\[(?:[^\[\]]|(?R))*\])'
    match = regex.search(pattern, text, regex.DOTALL)
    if match:
        return match.group(0)

    return text


def fix_unicode_text(text: str) -> str:
    fixed = ftfy.fix_text(text)
    fixed = fixed.encode("utf-8", errors="ignore").decode("utf-8")
    fixed = fixed.replace("\ufffd", " ")
    return fixed


def debug_print_trace():
    if settings.DEBUG:
        print(traceback.format_exc())


def exact_deduplicate(topics):
    result = []
    for line in topics:
        if not isinstance(line, (str, dict)):
            continue

        if isinstance(line, dict):
            topic = line["topic"]
        else:
            topic = line

        if topic not in result:
            result.append(line)
    return result

app/util.py
VikParuchuri/textbook_quality
import math
from typing import Optional, Dict
import argparse
import asyncio

from sentence_transformers import SentenceTransformer
from tqdm import tqdm

from app.db.session import get_session
from app.db.tables import * # Needed to avoid errors with table imports
from app.course.tasks import create_course_concepts, create_course_outline, query_course_context
from app.course.models import load_cached_course, Course
from app.lesson.tasks import generate_lesson
from app.lesson.output import render_components_to_output_markdown
from app.llm.generators.outline import renumber_outline
from app.settings import settings
import json
import os
import random
import ray

from app.util import debug_print_trace, exact_deduplicate


async def save_course(course: Course):
    async with get_session() as db:
        db.add(course)
        await db.commit()


def get_json_data_from_course(course: Course, extended_fields=False):
    json_data = {
        "topic": course.topic,
        "model": settings.LLM_TYPE,
        "concepts": course.concepts,
        "outline": course.outline,
        "markdown": course.markdown,
    }

    if extended_fields:
        if isinstance(course.components[0], str):
            json_data["components"] = course.components
        else:
            json_data["components"] = [c.json() for c in course.components]

        if course.context is None or len(course.context) == 0 or isinstance(course.context[0], str):
            json_data["context"] = course.context
        else:
            json_data["context"] = [v.json() for v in course.context]
        json_data["queries"] = course.queries
    return json.dumps(json_data)


async def generate_single_course(model, course_data: Dict | str, revision=1, outline_items=12, cache_only=False):
    components = ["exercise", "example"]

    outline = None
    queries = None
    concepts = []
    if isinstance(course_data, dict):
        course_name = course_data["topic"]
        outline = course_data["outline"]
        queries = course_data["queries"]
    else:
        course_name = course_data

    course = await load_cached_course(settings.LLM_TYPE, course_name, revision)
    if course is not None:
        await asyncio.sleep(0.01) # small sleep to avoid excess db load
        return course

    if cache_only:
        return None

    if not outline:
        # Only generate outline if one was not passed in
        concepts = await create_course_concepts(course_name, revision)
        if concepts is None:
            return

        outline, queries = await create_course_outline(course_name, concepts, outline_items, revision)

        if outline is None:
            return

        # Remove the intro if it exists
        if "intro" in outline[0].lower():
            outline = outline[1:]
            # Remove sections of intro
            while outline[0].startswith("1."):
                outline = outline[1:]
            outline = renumber_outline(outline)

    context = None
    if queries is not None:
        try:
            # Up to one retrieved passage per outline item
            # Remove numbers from outline for use in retrieval
            context_outline = [item.split(" ", 1)[-1] for item in outline]
            context = await query_course_context(model, queries, context_outline)
        except Exception as e:
            debug_print_trace()
            print(f"Error generating context for {course_name}: {e}")

    components = await generate_lesson(course_name, components, outline, revision, research_notes=context)
    if components is None:
        return

    md = render_components_to_output_markdown(components)

    course = Course(
        topic=course_name,
        model=settings.LLM_TYPE,
        outline=outline,
        concepts=concepts,
        markdown=md,
        components=components,
        context=context,
        queries=queries if queries is not None else [],
        version=revision
    )
    await save_course(course)

    return course


async def _process_course(model, topic, args):
    try:
        return await generate_single_course(model, topic, revision=args.revision, cache_only=args.cache_only)
    except Exception as e:
        debug_print_trace()
        print(f"Unhandled error generating course: {e}")


async def _process_courses(model, courses, args):
    processes = [_process_course(model, course, args) for course in courses]
    return await asyncio.gather(*processes)


@ray.remote(num_cpus=settings.RAY_CORES_PER_WORKER)
def process_courses(model, courses, args):
    try:
        return asyncio.run(_process_courses(model, courses, args))
    except Exception as e:
        debug_print_trace()
        print(f"Unhandled error generating courses: {e}")


@ray.remote(num_cpus=settings.RAY_CORES_PER_WORKER)
def process_course(model, course, args):
    try:
        return asyncio.run(_process_course(model, course, args))
    except Exception as e:
        debug_print_trace()
        print(f"Unhandled error generating course: {e}")


def load_topics(in_file: str):
    with open(os.path.join(settings.DATA_DIR, in_file)) as f:
        if in_file.endswith(".json"):
            topics = json.load(f)
        elif in_file.endswith(".jsonl"):
            lines = list(f)
            topics = []
            for line in lines:
                topics.append(json.loads(line))
        else:
            raise Exception(f"Unknown file type for {in_file}")

    random.seed(1)
    random.shuffle(topics)

    return topics


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Given a topic file, generate synthetic books.")
    parser.add_argument("in_file", help="Input filename (flat json list of topics, or jsonl file with dictionaries with keys topic, queries, and outline).  One file or a comma-separated list of files.")
    parser.add_argument("out_file", help="Output filename (jsonl)")
    parser.add_argument("--max", type=int, default=None, help="Maximum number of courses to generate")
    parser.add_argument("--workers", type=int, default=5, help="Number of workers to use")
    parser.add_argument("--extended-fields", action="store_true", default=False, help="Include extended fields in output")
    parser.add_argument("--revision", type=int, default=1, help="Revision number for the course.  Change this to avoid hitting cache if you want to regenerate a course.")
    parser.add_argument("--cache-only", action="store_true", default=False, help="Only use the cache, don't generate any new courses")

    args = parser.parse_args()

    # Load in topics, limit to max if needed
    # Also shuffle randomly (with a seed)
    topics = []
    in_files = args.in_file.split(",")
    for in_file in in_files:
        topics += load_topics(in_file.strip())

    if args.max is not None:
        topics = topics[:args.max]

    # Everything is cached, so exact duplicates will result in the same output
    topics = exact_deduplicate(topics)

    total_processes = math.ceil(args.workers * settings.RAY_CORES_PER_WORKER)
    func = process_course

    if settings.THREADS_PER_WORKER > 1:
        # group topics into batches of settings.THREADS_PER_WORKER
        topics = [topics[i:i + settings.THREADS_PER_WORKER] for i in
                          range(0, len(topics), settings.THREADS_PER_WORKER)]
        total_processes = math.ceil(total_processes / settings.THREADS_PER_WORKER)
        func = process_courses

    ray.init(
        num_cpus=total_processes,
        storage=settings.RAY_CACHE_PATH,
        _temp_dir=settings.RAY_CACHE_PATH,
        dashboard_host=settings.RAY_DASHBOARD_HOST
    )

    model = SentenceTransformer("TaylorAI/gte-tiny")
    model_ref = ray.put(model)

    print(f"Generating {len(topics)} course batches with {total_processes} processes from filename(s) {in_files}")
    futures = [func.remote(model_ref, batch, args) for batch in topics]

    # Run all ray tasks
    courses = []
    progress_bar = tqdm(total=len(futures))
    while len(futures) > 0:
        finished, futures = ray.wait(
            futures, timeout=7.0
        )
        course = ray.get(finished)
        courses.extend(course)
        progress_bar.update(len(course))

    if settings.THREADS_PER_WORKER > 1:
        # Flatten courses list
        courses = [course for batch in courses for course in batch]

    course_count = 0
    with open(os.path.join(settings.DATA_DIR, args.out_file), "w+") as f:
        for course in courses:
            # Filter out courses that didn't generate properly
            if course is None or isinstance(course, Exception):
                continue

            if course.markdown is None or len(course.markdown) == 0:
                continue

            course_count += 1
            json_data = get_json_data_from_course(course, extended_fields=args.extended_fields)
            f.write(json_data + '\n')
    print(f"Generated {course_count} courses")
    ray.shutdown()




book_generator.py
VikParuchuri/textbook_quality
import asyncio
import os, sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from sqlalchemy import delete
from sqlmodel import select

from app.course.models import Course
from app.db.session import get_session
from app.llm.models import Prompt

import argparse


async def clear_model_data(model: str, dry_run=False):
    async with get_session() as db:
        query = await db.exec(select(Prompt).where(Prompt.model == model))
        prompts = query.all()
        print(f"Found {len(prompts)} prompts for model {model}.")
        if not dry_run:
            await db.exec(delete(Prompt).where(Prompt.model == model))

        query = await db.exec(select(Course).where(Course.model == model))
        courses = query.all()
        print(f"Found {len(courses)} courses for model {model}.")
        if not dry_run:
            await db.exec(delete(Course).where(Course.model == model))
            await db.commit()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Remove all data related to a specific model.")
    parser.add_argument("model", help="model name")
    parser.add_argument("--dry-run", action="store_true", help="Don't actually delete anything.")
    args = parser.parse_args()

    asyncio.run(clear_model_data(args.model, args.dry_run))




scripts/clear_model_data.py
VikParuchuri/textbook_quality
import os, sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import argparse
from app.settings import settings
import json
import datasets
from transformers import AutoTokenizer


tokenizer = AutoTokenizer.from_pretrained("vikp/code_llama_7b_hf")


def book_stats(filename: str):
    courses = []

    with open(os.path.join(settings.DATA_DIR, filename)) as f:
        lines = list(f)
        for line in lines:
            courses.append(json.loads(line))

    data = datasets.Dataset.from_list(courses)

    tokens = tokenizer(data["markdown"])
    total = 0

    for i in range(len(data)):
        total += len(tokens[i])

    print(f"Total books: {len(data)}")
    print(f"Total tokens: {total}")
    print(f"Tokens per textbook: {total / len(data)}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Parse and print stats about generated textbooks.")
    parser.add_argument("in_file", help="Input filename (flat json list)")
    args = parser.parse_args()

    book_stats(args.in_file)




scripts/parse_books.py
VikParuchuri/textbook_quality
from invoke import task
from app.settings import settings

DEV_ENV = {
    "DATABASE_URL": settings.DATABASE_URL,
    "PYTHONPATH": "."
}

@task
def generate_migrations(ctx):
    ctx.run("alembic revision --autogenerate", env=DEV_ENV)


@task
def migrate_dev(ctx):
    ctx.run("alembic upgrade head", env=DEV_ENV)

@task
def downgrade_dev(ctx):
    ctx.run("alembic downgrade -1", env=DEV_ENV)


@task
def clear_dev_db(ctx):
    ctx.run('psql postgres -c "drop database textbook;"')


@task
def create_dev_db(ctx):
    ctx.run('psql postgres -c "create database textbook;"')


@task
def reset_dev_db(ctx):
    clear_dev_db(ctx)
    create_dev_db(ctx)
    migrate_dev(ctx)


@task
def reset_migrations(ctx):
    ctx.run("rm -rf alembic/versions/*")
    generate_migrations(ctx)
    migrate_dev(ctx)


@task
def lint(ctx):
    ctx.run("pylint --load-plugins pylint_pydantic app")


@task
def black(ctx):
    ctx.run(
        "autoflake --imports=sqlmodel,sqlalchemy,typing --recursive --in-place app --exclude=__init__.py,tables.py,main.py,db.py,alembic,data,cache")
    ctx.run("isort app --profile black")
    ctx.run("black app --exclude tables.py,main.py,db.py,__init__.py,alembic,data,cache")

tasks.py
VikParuchuri/textbook_quality
import argparse

from tqdm.contrib.concurrent import process_map

from app.llm.exceptions import GenerationError, InvalidRequestError
from app.settings import settings
import json
import os
from app.llm.generators.toc import generate_tocs
import asyncio

from app.util import debug_print_trace


async def generate_tocs_async(topic, toc):
    final = await generate_tocs(topic, toc)
    return final


def generate_tocs_sync(topic, toc):
    try:
        return asyncio.run(generate_tocs_async(topic, toc))
    except (GenerationError, InvalidRequestError):
        debug_print_trace()
        print(f"Error generating toc for {topic}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Clean up real ToCs from books.")
    parser.add_argument("in_file", help="Input tocs and titles in jsonl format.  Should be a list of json objects with the keys toc (str) and title (str).", type=str)
    parser.add_argument("out_file", help="Output filename in jsonl format (list of json objects)")
    parser.add_argument("--max", type=int, default=None, help="Maximum number of tocs to generate")
    parser.add_argument("--workers", type=int, default=5, help="Number of workers to use")
    args = parser.parse_args()

    topics = []
    with open(os.path.join(settings.DATA_DIR, args.in_file)) as f:
        lines = list(f)

    for line in lines:
        topics.append(json.loads(line))

    if args.max is not None:
        topics = topics[:args.max]

    new_tocs = process_map(generate_tocs_sync, [t["title"] for t in topics], [t["toc"] for t in topics], max_workers=args.workers, chunksize=1)
    new_tocs = [t for t in new_tocs if t is not None]

    json_tocs = [t.dict() for t in new_tocs]
    with open(os.path.join(settings.DATA_DIR, args.out_file), "w+") as f:
        for line in json_tocs:
            f.write(json.dumps(line))
            f.write("\n")

    print(f"Generated {len(json_tocs)} tocs")


toc_cleaner.py
VikParuchuri/textbook_quality
import traceback
from copy import copy
import argparse

from app.settings import settings
import json
import os
from app.llm.generators.topic import generate_topic, generate_specific_topic
from app.course.embeddings import dedup_list
import asyncio
from tqdm import tqdm

from app.util import debug_print_trace


def load_processed_titles(file_name: str):
    with open(os.path.join(settings.DATA_DIR, file_name)) as f:
        titles = json.load(f)
    return titles


async def generate_topics(title):
    topics = await generate_topic(title)
    return topics

async def generate_specific_topics(title, domain=None):
    topics = await generate_specific_topic(title, domain)
    return topics


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Augment book titles with synthetic topics, and dedup.")
    parser.add_argument("in_file", help="Input filename (flat json list)")
    parser.add_argument("out_file", help="Output filename (flat json list)")
    parser.add_argument("--domain", help="Specific domain for the topics", default=None, type=str)
    parser.add_argument("--max", type=int, default=None, help="Maximum number of courses to generate")
    args = parser.parse_args()

    titles = load_processed_titles(args.in_file)
    # Filter titles to max if needed
    if args.max:
        titles = titles[:args.max]

    topics_from_titles = copy(titles)
    for title in tqdm(titles):
        try:
            topics = asyncio.run(generate_topics(title))
            topics_from_titles.extend(topics)
        except Exception as e:
            debug_print_trace()
            print(f"Error generating topic: {e}")

    topics_from_titles = dedup_list(topics_from_titles)
    print(len(topics_from_titles))

    all_topics = copy(topics_from_titles)
    for topic in tqdm(topics_from_titles):
        try:
            topics = asyncio.run(generate_specific_topics(topic, domain=args.domain))
            all_topics.extend(topics)
        except Exception as e:
            debug_print_trace()
            print(f"Error generating specific topic: {e}")

    all_topics = dedup_list(all_topics)

    with open(os.path.join(settings.DATA_DIR, args.out_file), "w+") as f:
        json.dump(all_topics, f, indent=2)

    print(len(all_topics))


topic_augmentor.py
VikParuchuri/textbook_quality
import traceback
import argparse

from app.llm.exceptions import GenerationError
from app.settings import settings
import json
import os
from app.llm.generators.title import generate_title
from app.course.embeddings import dedup_list
import asyncio
from tqdm import tqdm

from app.util import debug_print_trace

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate book titles from a given subject.")
    parser.add_argument("subject", help="Input subject", type=str)
    parser.add_argument("out_file", help="Output filename (flat json list)")
    parser.add_argument("--iterations", default=10, type=int, help="Number of times to generate titles")
    args = parser.parse_args()

    all_topics = []
    for i in tqdm(range(args.iterations)):
        try:
            topics = asyncio.run(generate_title(args.subject))
            all_topics.extend(topics)
        except GenerationError as e:
            debug_print_trace()
            print(f"Error generating titles {e}")

    all_topics = dedup_list(all_topics, score_thresh=.95)
    with open(os.path.join(settings.DATA_DIR, args.out_file), "w+") as f:
        json.dump(all_topics, f, indent=2)

    print(len(all_topics))


topic_generator.py
VikParuchuri/textbook_quality